From ba0bdc5fe38ad632759e4086ecede7df2d96c922 Mon Sep 17 00:00:00 2001
From: Wei <xnwei@amazon.com>
Date: Thu, 21 Sep 2017 11:26:10 -0700
Subject: [PATCH] Addressed comments

---
 .../lib/worker/KinesisClientLibConfiguration.java        | 6 ++++++
 .../clientlibrary/lib/worker/RecordsFetcherFactory.java  | 2 ++
 .../lib/worker/SimpleRecordsFetcherFactory.java          | 9 ++++++++-
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java
index e62be1ee..2c881de3 100644
--- a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java
+++ b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java
@@ -1294,6 +1294,12 @@ public class KinesisClientLibConfiguration {
         return this;
     }
 
+    public KinesisClientLibConfiguration withMaxRecordsCount(final int maxRecordsCount) {
+        checkIsValuePositive("maxRecordsCount", maxRecordsCount);
+        recordsFetcherFactory.setMaxRecordsCount(maxRecordsCount);
+        return this;
+    }
+
     /**
      * @param timeoutInSeconds The timeout in seconds to wait for the MultiLangProtocol to wait for
      */
diff --git a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/RecordsFetcherFactory.java b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/RecordsFetcherFactory.java
index 86e78ac4..cdd80e49 100644
--- a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/RecordsFetcherFactory.java
+++ b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/RecordsFetcherFactory.java
@@ -32,6 +32,8 @@ public interface RecordsFetcherFactory {
 
     void setMaxByteSize(int maxByteSize);
 
+    void setMaxRecordsCount(int maxRecordsCount);
+
     void setDataFetchingStrategy(DataFetchingStrategy dataFetchingStrategy);
 
 }
diff --git a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/SimpleRecordsFetcherFactory.java b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/SimpleRecordsFetcherFactory.java
index 465110c1..cf0c0a2c 100644
--- a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/SimpleRecordsFetcherFactory.java
+++ b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/SimpleRecordsFetcherFactory.java
@@ -14,7 +14,6 @@
  */
 package com.amazonaws.services.kinesis.clientlibrary.lib.worker;
 
-import lombok.Setter;
 import lombok.extern.apachecommons.CommonsLog;
 
 import java.util.concurrent.Executors;
@@ -48,14 +47,22 @@ public class SimpleRecordsFetcherFactory implements RecordsFetcherFactory {
         }
     }
 
+    @Override
     public void setMaxSize(int maxSize){
         this.maxSize = maxSize;
     }
 
+    @Override
     public void setMaxByteSize(int maxByteSize){
         this.maxByteSize = maxByteSize;
     }
 
+    @Override
+    public void setMaxRecordsCount(int maxRecordsCount) {
+        this.maxRecordsCount = maxRecordsCount;
+    }
+
+    @Override
     public void setDataFetchingStrategy(DataFetchingStrategy dataFetchingStrategy){
         this.dataFetchingStrategy = dataFetchingStrategy;
     }