Always return at least one record when exceeds buffer size

src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTask.java

@@ -148,7 +148,7 @@ public List<SourceRecord> poll() throws InterruptedException {
         while (running.get()) {
             fillRecords(records, topic);            
             if (records.isEmpty() || System.currentTimeMillis() > maxWaitTime) {
-                logger.debug("Sending {} documents.", records.size());
+                logger.info("Sending {} documents.", records.size());
                 break;
             }
         }
@@ -179,9 +179,13 @@ private void fillRecords(List<SourceRecord> records, String topic) throws Interr
                 }
 
                 // Get the latest token and record as offset
+                // TODO: The continuationToken here is picked from any lease with owner, so maybe a little bit random
+                // change to show the continuationToken for the leases processed by the current worker
                 Map<String, Object> sourceOffset = singletonMap(OFFSET_KEY, getContinuationToken());
-                logger.debug("Latest offset is {}.", sourceOffset.get(OFFSET_KEY));
 
+                if (logger.isDebugEnabled()) {
+                    logger.debug("Latest offset is {}.", sourceOffset.get(OFFSET_KEY));
+                }
                 // Convert JSON to Kafka Connect struct and JSON schema
                 SchemaAndValue schemaAndValue = jsonToStruct.recordToSchemaAndValue(node);
 
@@ -192,14 +196,19 @@ private void fillRecords(List<SourceRecord> records, String topic) throws Interr
 
                 bufferSize -= sourceRecord.value().toString().getBytes().length;
 
-                // If the buffer Size exceeds then do not remove the node .
-                if (bufferSize <= 0) {
+                // Add the item to buffer if either conditions met:
+                // it is the first record, or adding this record does not exceed the buffer size
+                if (records.size() == 0 || bufferSize >= 0) {
+                    records.add(sourceRecord);
+                    count++;
+                } else {
+                    // If the buffer Size exceeds then do not remove the node.
+                    if (logger.isDebugEnabled()) {
+                        logger.debug("Adding record back to the queue since adding it exceeds the allowed buffer size {}", config.getTaskBufferSize());
+                    }
                     this.queue.add(node);
                     break;
                 }
-
-                records.add(sourceRecord);
-                count++;
             } catch (Exception e) {
                 logger.error("Failed to fill Source Records for Topic {}", topic);
                 throw e;

src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTaskTest.java

@@ -159,7 +159,7 @@ public void testSmallBufferSize() throws InterruptedException, JsonProcessingExc
         }).start();
 
         List<SourceRecord>  result=testTask.poll();
-        Assert.assertEquals(0, result.size());
+        Assert.assertEquals(1, result.size());
     }