@ProcessElement
  public ProcessContinuation processElement(
      @Element byte[] element,
      RestrictionTracker<OffsetRange, Long> tracker,
      WatermarkEstimator<Instant> watermarkEstimator,
      OutputReceiver<V> receiver) {

    if (tracker.currentRestriction() != null) {
      LOG.info(
          "Start processing element. Restriction = {}", tracker.currentRestriction().toString());
    }
    SparkConsumer<V> sparkConsumer;
    Receiver<V> sparkReceiver;
    try {
      sparkReceiver = sparkReceiverBuilder.build();
    } catch (Exception e) {
      LOG.error("Can not build Spark Receiver", e);
      throw new IllegalStateException("Spark Receiver was not built!");
    }
    LOG.debug("Restriction {}", tracker.currentRestriction().toString());
    sparkConsumer = new SparkConsumerWithOffset<>(tracker.currentRestriction().getFrom());
    sparkConsumer.start(sparkReceiver);

    Long recordsProcessed = 0L;
    while (true) {
      LOG.debug("Start polling records");
      try {
        TimeUnit.SECONDS.sleep(startPollTimeoutSec);
      } catch (InterruptedException e) {
        LOG.error("SparkReceiver was interrupted before polling started", e);
        throw new IllegalStateException("Spark Receiver was interrupted before polling started");
      }
      if (!sparkConsumer.hasRecords()) {
        LOG.debug("No records left");
        ((HasOffset) sparkReceiver).setCheckpoint(recordsProcessed);
        sparkConsumer.stop();
        tracker.checkDone();
        if (pullFrequencySec != 0L) {
          LOG.debug("Waiting to poll for new records...");
          try {
            TimeUnit.SECONDS.sleep(pullFrequencySec);
          } catch (InterruptedException e) {
            LOG.error("SparkReceiver was interrupted while waiting to poll new records", e);
            throw new IllegalStateException(
                "Spark Receiver was interrupted while waiting to poll new records");
          }
        }
        OffsetRange currentRestriction = tracker.currentRestriction();
        if (currentRestriction != null
            && currentRestriction.getFrom() == currentRestriction.getTo()) {
          LOG.info("Stop for empty restriction: {}", currentRestriction);
          return ProcessContinuation.stop();
        } else {
          LOG.info("Resume for restriction: {}", currentRestriction);
          return ProcessContinuation.resume();
        }
      }
      while (sparkConsumer.hasRecords()) {
        V record = sparkConsumer.poll();
        if (record != null) {
          Long offset = getOffsetFn.apply(record);
          if (!tracker.tryClaim(offset)) {
            ((HasOffset) sparkReceiver).setCheckpoint(recordsProcessed);
            sparkConsumer.stop();
            LOG.info("Stop for restriction: {}", tracker.currentRestriction());
            return ProcessContinuation.stop();
          }
          Instant currentTimeStamp = getTimestampFn.apply(record);
          recordsProcessed++;
          ((ManualWatermarkEstimator<Instant>) watermarkEstimator).setWatermark(currentTimeStamp);
          receiver.outputWithTimestamp(record, currentTimeStamp);
        }
      }
    }
  }
@Test
  public void testProcessElement() {
    MockOutputReceiver receiver = new MockOutputReceiver();
    DoFn.ProcessContinuation result =
        dofnInstance.processElement(
            TEST_ELEMENT,
            dofnInstance.restrictionTracker(
                TEST_ELEMENT, dofnInstance.initialRestriction(TEST_ELEMENT)),
            mockWatermarkEstimator,
            receiver);
    assertEquals(DoFn.ProcessContinuation.resume(), result);
    assertEquals(
        createExpectedRecords(CustomReceiverWithOffset.RECORDS_COUNT), receiver.getOutputs());
  }