address comments

Signed-off-by: Yupeng Fu <[email protected]>

Author: yupeng9

plugins/ingestion-kafka/src/main/java/org/opensearch/plugin/kafka/KafkaPartitionConsumer.java

@@ -136,6 +136,19 @@ public List<ReadResult<KafkaOffset, KafkaMessage>> readNext(
         return records;
     }
 
+    @Override
+    public List<ReadResult<KafkaOffset, KafkaMessage>> readNext(long maxMessages, int timeoutMillis) throws TimeoutException {
+        List<ReadResult<KafkaOffset, KafkaMessage>> records = AccessController.doPrivileged(
+            (PrivilegedAction<List<ReadResult<KafkaOffset, KafkaMessage>>>) () -> fetch(
+                lastFetchedOffset,
+                false,
+                maxMessages,
+                timeoutMillis
+            )
+        );
+        return records;
+    }
+
     @Override
     public IngestionShardPointer earliestPointer() {
         long startOffset = AccessController.doPrivileged(

plugins/ingestion-kinesis/src/internalClusterTest/java/org/opensearch/plugin/kinesis/KinesisIngestionBaseIT.java

@@ -91,6 +91,7 @@ private void stopKinesis() {
 
         if (localstack != null) {
             localstack.stop();
+            localstack = null;
         }
     }
 

plugins/ingestion-kinesis/src/main/java/org/opensearch/plugin/kinesis/KinesisShardConsumer.java

@@ -49,7 +49,7 @@ public class KinesisShardConsumer implements IngestionShardConsumer<SequenceNumb
      */
     ;
     private KinesisClient kinesisClient;
-    private String lastFetchedSequenceNumber = "";
+    private String lastShardIterator;
     final String clientId;
     final String kinesisShardId;
     final int shardId;
@@ -124,6 +124,7 @@ public List<ReadResult<SequenceNumber, KinesisMessage>> readNext(
         int timeoutMillis
     ) throws TimeoutException {
         List<ReadResult<SequenceNumber, KinesisMessage>> records = fetch(
+            null,
             sequenceNumber.getSequenceNumber(),
             includeStart,
             maxMessages,
@@ -132,6 +133,14 @@ public List<ReadResult<SequenceNumber, KinesisMessage>> readNext(
         return records;
     }
 
+    @Override
+    public List<ReadResult<SequenceNumber, KinesisMessage>> readNext(long maxMessages, int timeoutMillis) throws TimeoutException {
+        if (lastShardIterator == null) {
+            throw new IllegalStateException("No shard iterator available");
+        }
+        return fetch(lastShardIterator, null, false, maxMessages, timeoutMillis);
+    }
+
     @Override
     public IngestionShardPointer earliestPointer() {
         return getSequenceNumber(ShardIteratorType.TRIM_HORIZON, null, 0);
@@ -142,40 +151,50 @@ public IngestionShardPointer latestPointer() {
         return getSequenceNumber(ShardIteratorType.LATEST, null, 0);
     }
 
-    private List<Record> fetchRecords(ShardIteratorType shardIteratorType, String startingSequenceNumber, long timestampMillis, int limit) {
-        // Get a shard iterator AFTER the given sequence number
-        GetShardIteratorRequest.Builder builder = GetShardIteratorRequest.builder()
-            .streamName(config.getStream())
-            .shardId(kinesisShardId)
-            .shardIteratorType(shardIteratorType);
+    private List<Record> fetchRecords(
+        String shardIterator,
+        ShardIteratorType shardIteratorType,
+        String startingSequenceNumber,
+        long timestampMillis,
+        int limit
+    ) {
+        String shardIteratorToUse = shardIterator;
 
-        if (startingSequenceNumber != null) {
-            builder = builder.startingSequenceNumber(startingSequenceNumber);
-        }
+        if (shardIterator == null) {
+            // fetch the shard iterator
+            GetShardIteratorRequest.Builder builder = GetShardIteratorRequest.builder()
+                .streamName(config.getStream())
+                .shardId(kinesisShardId)
+                .shardIteratorType(shardIteratorType);
+
+            if (startingSequenceNumber != null) {
+                builder = builder.startingSequenceNumber(startingSequenceNumber);
+            }
 
-        if (timestampMillis != 0) {
-            builder = builder.timestamp(Instant.ofEpochMilli(timestampMillis));
-        }
+            if (timestampMillis != 0) {
+                builder = builder.timestamp(Instant.ofEpochMilli(timestampMillis));
+            }
 
-        GetShardIteratorRequest shardIteratorRequest = builder.build();
+            GetShardIteratorRequest shardIteratorRequest = builder.build();
 
-        GetShardIteratorResponse shardIteratorResponse = kinesisClient.getShardIterator(shardIteratorRequest);
-        String shardIterator = shardIteratorResponse.shardIterator();
+            GetShardIteratorResponse shardIteratorResponse = kinesisClient.getShardIterator(shardIteratorRequest);
+            shardIteratorToUse = shardIteratorResponse.shardIterator();
+        }
 
-        if (shardIterator == null) {
+        if (shardIteratorToUse == null) {
             return new ArrayList<>();
         }
 
         // Fetch the next records
-        GetRecordsRequest recordsRequest = GetRecordsRequest.builder().shardIterator(shardIterator).limit(limit).build();
-
+        GetRecordsRequest recordsRequest = GetRecordsRequest.builder().shardIterator(shardIteratorToUse).limit(limit).build();
         GetRecordsResponse recordsResponse = kinesisClient.getRecords(recordsRequest);
+        lastShardIterator = recordsResponse.nextShardIterator();
         List<Record> records = recordsResponse.records();
         return records;
     }
 
     private SequenceNumber getSequenceNumber(ShardIteratorType shardIteratorType, String startingSequenceNumber, long timestampMillis) {
-        List<Record> records = fetchRecords(shardIteratorType, startingSequenceNumber, timestampMillis, 1);
+        List<Record> records = fetchRecords(null, shardIteratorType, startingSequenceNumber, timestampMillis, 1);
 
         if (!records.isEmpty()) {
             Record nextRecord = records.get(0);
@@ -197,6 +216,7 @@ public IngestionShardPointer pointerFromOffset(String offset) {
     }
 
     private synchronized List<ReadResult<SequenceNumber, KinesisMessage>> fetch(
+        String shardIterator,
         String sequenceNumber,
         boolean includeStart,
         long maxMessages,
@@ -208,7 +228,7 @@ private synchronized List<ReadResult<SequenceNumber, KinesisMessage>> fetch(
 
         ShardIteratorType iteratorType = includeStart ? ShardIteratorType.AT_SEQUENCE_NUMBER : ShardIteratorType.AFTER_SEQUENCE_NUMBER;
 
-        List<Record> records = fetchRecords(iteratorType, sequenceNumber, 0, (int) limit);
+        List<Record> records = fetchRecords(shardIterator, iteratorType, sequenceNumber, 0, (int) limit);
 
         List<ReadResult<SequenceNumber, KinesisMessage>> results = new ArrayList<>();
 

server/src/main/java/org/opensearch/index/IngestionShardConsumer.java

@@ -69,6 +69,16 @@ public M getMessage() {
     List<ReadResult<T, M>> readNext(T pointer, boolean includeStart, long maxMessages, int timeoutMillis)
         throws java.util.concurrent.TimeoutException;
 
+    /**
+     * Read the next set of messages from the source using the previous pointer. An exception is thrown if no previous pointer is available.
+     * This method is used as an optimization for consecutive reads.
+     * @param maxMessages the maximum number of messages to read, or -1 for no limit
+     * @param timeoutMillis the maximum time to wait for messages
+     * @return a list of messages read from the source
+     * @throws java.util.concurrent.TimeoutException
+     */
+    List<ReadResult<T, M>> readNext(long maxMessages, int timeoutMillis) throws java.util.concurrent.TimeoutException;
+
     /**
      * @return the earliest pointer in the shard
      */

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

@@ -149,9 +149,6 @@ protected void startPoll() {
         }
         logger.info("Starting poller for shard {}", consumer.getShardId());
 
-        // track the last record successfully written to the blocking queue
-        IngestionShardPointer lastSuccessfulPointer = null;
-
         while (true) {
             try {
                 if (closed) {
@@ -204,7 +201,7 @@ protected void startPoll() {
                 if (includeBatchStartPointer) {
                     results = consumer.readNext(batchStartPointer, true, MAX_POLL_SIZE, POLL_TIMEOUT);
                 } else {
-                    results = consumer.readNext(lastSuccessfulPointer, false, MAX_POLL_SIZE, POLL_TIMEOUT);
+                    results = consumer.readNext(MAX_POLL_SIZE, POLL_TIMEOUT);
                 }
 
                 if (results.isEmpty()) {
@@ -221,7 +218,6 @@ protected void startPoll() {
                         batchStartPointer = result.getPointer();
                         firstInBatch = false;
                     }
-                    lastSuccessfulPointer = result.getPointer();
 
                     // check if the message is already processed
                     if (isProcessed(result.getPointer())) {