Use shard pointer tracked by writer for recovery

Signed-off-by: Varun Bharadwaj <[email protected]>

Author: varunbharadwaj

server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java

@@ -51,7 +51,7 @@ public class DefaultStreamPoller implements StreamPoller {
     private ExecutorService processorThread;
 
     // start of the batch, inclusive
-    private IngestionShardPointer batchStartPointer;
+    private IngestionShardPointer initialBatchStartPointer;
     private boolean includeBatchStartPointer = false;
 
     private ResetState resetState;
@@ -104,7 +104,7 @@ public DefaultStreamPoller(
         this.consumer = Objects.requireNonNull(consumer);
         this.resetState = resetState;
         this.resetValue = resetValue;
-        this.batchStartPointer = startPointer;
+        this.initialBatchStartPointer = startPointer;
         this.state = initialState;
         this.persistedPointers = persistedPointers;
         if (!this.persistedPointers.isEmpty()) {
@@ -169,23 +169,23 @@ protected void startPoll() {
                 if (resetState != ResetState.NONE) {
                     switch (resetState) {
                         case EARLIEST:
-                            batchStartPointer = consumer.earliestPointer();
-                            logger.info("Resetting offset by seeking to earliest offset {}", batchStartPointer.asString());
+                            initialBatchStartPointer = consumer.earliestPointer();
+                            logger.info("Resetting offset by seeking to earliest offset {}", initialBatchStartPointer.asString());
                             break;
                         case LATEST:
-                            batchStartPointer = consumer.latestPointer();
-                            logger.info("Resetting offset by seeking to latest offset {}", batchStartPointer.asString());
+                            initialBatchStartPointer = consumer.latestPointer();
+                            logger.info("Resetting offset by seeking to latest offset {}", initialBatchStartPointer.asString());
                             break;
                         case REWIND_BY_OFFSET:
-                            batchStartPointer = consumer.pointerFromOffset(resetValue);
-                            logger.info("Resetting offset by seeking to offset {}", batchStartPointer.asString());
+                            initialBatchStartPointer = consumer.pointerFromOffset(resetValue);
+                            logger.info("Resetting offset by seeking to offset {}", initialBatchStartPointer.asString());
                             break;
                         case REWIND_BY_TIMESTAMP:
-                            batchStartPointer = consumer.pointerFromTimestampMillis(Long.parseLong(resetValue));
+                            initialBatchStartPointer = consumer.pointerFromTimestampMillis(Long.parseLong(resetValue));
                             logger.info(
                                 "Resetting offset by seeking to timestamp {}, corresponding offset {}",
                                 resetValue,
-                                batchStartPointer.asString()
+                                initialBatchStartPointer.asString()
                             );
                             break;
                     }
@@ -208,7 +208,8 @@ protected void startPoll() {
                 List<IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message>> results;
 
                 if (includeBatchStartPointer) {
-                    results = consumer.readNext(batchStartPointer, true, MAX_POLL_SIZE, POLL_TIMEOUT);
+                    results = consumer.readNext(initialBatchStartPointer, true, MAX_POLL_SIZE, POLL_TIMEOUT);
+                    includeBatchStartPointer = false;
                 } else {
                     results = consumer.readNext(MAX_POLL_SIZE, POLL_TIMEOUT);
                 }
@@ -219,38 +220,42 @@ protected void startPoll() {
                 }
 
                 state = State.PROCESSING;
-                // process the records
-                boolean firstInBatch = true;
-                for (IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message> result : results) {
-                    if (firstInBatch) {
-                        // update the batch start pointer to the next batch
-                        batchStartPointer = result.getPointer();
-                        firstInBatch = false;
-                    }
+                processRecords(results);
+            } catch (Throwable e) {
+                // Pause ingestion when an error is encountered while polling the streaming source.
+                // Currently we do not have a good way to skip past the failing messages.
+                // The user will have the option to manually update the offset and resume ingestion.
+                // todo: support retry?
+                logger.error("Pausing ingestion. Fatal error occurred in polling the shard {}: {}", consumer.getShardId(), e);
+                pause();
+            }
+        }
+    }
 
-                    // check if the message is already processed
-                    if (isProcessed(result.getPointer())) {
-                        logger.info("Skipping message with pointer {} as it is already processed", result.getPointer().asString());
-                        continue;
-                    }
-                    totalPolledCount.inc();
-                    blockingQueue.put(result);
-
-                    logger.debug(
-                        "Put message {} with pointer {} to the blocking queue",
-                        String.valueOf(result.getMessage().getPayload()),
-                        result.getPointer().asString()
-                    );
+    private void processRecords(List<IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message>> results) {
+        for (IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message> result : results) {
+            try {
+                // check if the message is already processed
+                if (isProcessed(result.getPointer())) {
+                    logger.info("Skipping message with pointer {} as it is already processed", result.getPointer().asString());
+                    continue;
                 }
-                // for future reads, we do not need to include the batch start pointer, and read from the last successful pointer.
-                includeBatchStartPointer = false;
+                totalPolledCount.inc();
+                blockingQueue.put(result);
+
+                logger.debug(
+                    "Put message {} with pointer {} to the blocking queue",
+                    String.valueOf(result.getMessage().getPayload()),
+                    result.getPointer().asString()
+                );
             } catch (Throwable e) {
-                logger.error("Error in polling the shard {}: {}", consumer.getShardId(), e);
+                logger.error("Error in processing a record. Shard {}, pointer {}: {}", consumer.getShardId(), result.getPointer(), e);
                 errorStrategy.handleError(e, IngestionErrorStrategy.ErrorStage.POLLING);
 
                 if (!errorStrategy.shouldIgnoreError(e, IngestionErrorStrategy.ErrorStage.POLLING)) {
                     // Blocking error encountered. Pause poller to stop processing remaining updates.
                     pause();
+                    break;
                 }
             }
         }
@@ -328,9 +333,16 @@ public boolean isClosed() {
         return closed;
     }
 
+    /**
+     * Returns the batch start pointer from where the poller can resume in case of shard recovery. The poller and
+     * processor are decoupled in this implementation, and hence the latest pointer tracked by the processor acts as the
+     * recovery/start point. In case the processor has not started tracking, then the initial batchStartPointer used by
+     * the poller acts as the start point.
+     */
     @Override
     public IngestionShardPointer getBatchStartPointer() {
-        return batchStartPointer;
+        IngestionShardPointer currentShardPointer = processorRunnable.getCurrentShardPointer();
+        return currentShardPointer == null ? initialBatchStartPointer : currentShardPointer;
     }
 
     @Override

server/src/main/java/org/opensearch/indices/pollingingest/MessageProcessorRunnable.java

@@ -13,6 +13,7 @@
 import org.apache.lucene.document.StoredField;
 import org.apache.lucene.index.Term;
 import org.opensearch.action.DocWriteRequest;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.lucene.uid.Versions;
 import org.opensearch.common.metrics.CounterMetric;
 import org.opensearch.common.util.RequestUtils;
@@ -59,6 +60,10 @@ public class MessageProcessorRunnable implements Runnable {
     private final MessageProcessor messageProcessor;
     private final CounterMetric stats = new CounterMetric();
 
+    // tracks the most recent pointer that is being processed
+    @Nullable
+    private volatile IngestionShardPointer currentShardPointer;
+
     /**
      * Constructor.
      *
@@ -274,6 +279,7 @@ public void run() {
             if (readResult != null) {
                 try {
                     stats.inc();
+                    currentShardPointer = readResult.getPointer();
                     messageProcessor.process(readResult.getMessage(), readResult.getPointer());
                     readResult = null;
                 } catch (Exception e) {
@@ -308,4 +314,9 @@ public IngestionErrorStrategy getErrorStrategy() {
     public void setErrorStrategy(IngestionErrorStrategy errorStrategy) {
         this.errorStrategy = errorStrategy;
     }
+
+    @Nullable
+    public IngestionShardPointer getCurrentShardPointer() {
+        return currentShardPointer;
+    }
 }

server/src/test/java/org/opensearch/indices/pollingingest/DefaultStreamPollerTests.java

@@ -374,4 +374,54 @@ public void testUpdateErrorStrategy() {
         assertTrue(poller.getErrorStrategy() instanceof BlockIngestionErrorStrategy);
         assertTrue(processorRunnable.getErrorStrategy() instanceof BlockIngestionErrorStrategy);
     }
+
+    public void testPersistedBatchStartPointer() throws TimeoutException, InterruptedException {
+        messages.add("{\"_id\":\"3\",\"_source\":{\"name\":\"bob\", \"age\": 24}}".getBytes(StandardCharsets.UTF_8));
+        messages.add("{\"_id\":\"4\",\"_source\":{\"name\":\"alice\", \"age\": 21}}".getBytes(StandardCharsets.UTF_8));
+        List<
+            IngestionShardConsumer.ReadResult<
+                FakeIngestionSource.FakeIngestionShardPointer,
+                FakeIngestionSource.FakeIngestionMessage>> readResultsBatch1 = fakeConsumer.readNext(
+                    fakeConsumer.earliestPointer(),
+                    true,
+                    2,
+                    100
+                );
+        List<
+            IngestionShardConsumer.ReadResult<
+                FakeIngestionSource.FakeIngestionShardPointer,
+                FakeIngestionSource.FakeIngestionMessage>> readResultsBatch2 = fakeConsumer.readNext(
+                    new FakeIngestionSource.FakeIngestionShardPointer(2),
+                    true,
+                    2,
+                    100
+                );
+
+        // This test publishes 4 messages, so use blocking queue of size 3. This ensures the poller is blocked when adding the 4th message
+        // for validation.
+        IngestionErrorStrategy errorStrategy = spy(new BlockIngestionErrorStrategy("ingestion_source"));
+        doThrow(new RuntimeException()).when(processor).process(any(), any());
+        processorRunnable = new MessageProcessorRunnable(new ArrayBlockingQueue<>(3), processor, errorStrategy);
+
+        IngestionShardConsumer mockConsumer = mock(IngestionShardConsumer.class);
+        when(mockConsumer.getShardId()).thenReturn(0);
+        when(mockConsumer.readNext(any(), anyBoolean(), anyLong(), anyInt())).thenReturn(readResultsBatch1);
+
+        when(mockConsumer.readNext(anyLong(), anyInt())).thenReturn(readResultsBatch2).thenReturn(Collections.emptyList());
+
+        poller = new DefaultStreamPoller(
+            new FakeIngestionSource.FakeIngestionShardPointer(0),
+            persistedPointers,
+            mockConsumer,
+            processorRunnable,
+            StreamPoller.ResetState.NONE,
+            "",
+            errorStrategy,
+            StreamPoller.State.NONE
+        );
+        poller.start();
+        Thread.sleep(sleepTime);
+
+        assertEquals(new FakeIngestionSource.FakeIngestionShardPointer(0), poller.getBatchStartPointer());
+    }
 }