pubsub: ensure atleast 1 message is pulled (#197)

Author: Yaminik1996

pubsub/pubsub-client/src/main/java/com/salesforce/multicloudj/pubsub/driver/AbstractSubscription.java

@@ -133,7 +133,7 @@ private QueueConfig() {
     protected final AtomicBoolean isShutdown = new AtomicBoolean(false);
 
     /** Current best estimate for how many messages to fetch in order to maintain the desired queue duration. */
-    private int runningBatchSize = 1;
+    private double runningBatchSize = 1;
 
     /** Timestamp (in nanoseconds) when the current throughput measurement window started. */
     private long throughputStart = 0L;
@@ -520,18 +520,19 @@ private int updateBatchSize() {
             double minSize = runningBatchSize * QueueConfig.MAX_SHRINK_FACTOR;
 
             if (newBatchSize > maxSize) {
-                runningBatchSize = (int) maxSize;
+                runningBatchSize = maxSize;
             } else if (newBatchSize < minSize) {
-                runningBatchSize = (int) minSize;
+                runningBatchSize = minSize;
             } else {
-                runningBatchSize = (int) newBatchSize;
+                runningBatchSize = newBatchSize;
             }
         }
 
         throughputStart = now;
         throughputCount = 0;
 
-        return Math.min(runningBatchSize, QueueConfig.MAX_BATCH_SIZE);
+        // Rounding up to the nearest integer to avoid pulling 0 messages
+        return (int) Math.ceil(Math.min(runningBatchSize, QueueConfig.MAX_BATCH_SIZE));
     }
 
     /**

pubsub/pubsub-client/src/test/java/com/salesforce/multicloudj/pubsub/client/AbstractPubsubIT.java

@@ -11,6 +11,7 @@
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Assumptions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Disabled;
@@ -26,6 +27,7 @@
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 public abstract class AbstractPubsubIT {
 
+    protected static final String AWS_PROVIDER_ID = "aws";
 
     public interface Harness extends AutoCloseable {
 
@@ -347,4 +349,60 @@ public void testGetAttributes() throws Exception {
             Assertions.assertNotNull(attributes.getTopic(), "Should have topic attribute");
         }
     }
+
+    @Test
+    @Timeout(30)
+    public void testMultipleSendReceiveWithoutBatch() throws Exception {
+        Assumptions.assumeFalse(AWS_PROVIDER_ID.equals(harness.getProviderId()));
+        try (AbstractTopic topic = harness.createTopicDriver();
+             AbstractSubscription subscription = harness.createSubscriptionDriver()) {
+
+            int numMessages = 5;
+            List<Message> sentMessages = new ArrayList<>();
+            
+            // Send messages one by one (not in batch)
+            for (int i = 0; i < numMessages; i++) {
+                Message message = Message.builder()
+                        .withBody(("non-batch-msg-" + i).getBytes())
+                        .withMetadata(Map.of("index", String.valueOf(i)))
+                        .build();
+                topic.send(message);
+                sentMessages.add(message);
+                TimeUnit.MILLISECONDS.sleep(100); // Small delay between sends
+            }
+
+            TimeUnit.MILLISECONDS.sleep(500); // Allow time for messages to be available
+
+            // Receive and ack messages one by one (not in batch)
+            List<Message> receivedMessages = new ArrayList<>();
+            
+            while (receivedMessages.size() < numMessages) {
+                try {
+                    Message received = subscription.receive();
+                    if (received != null && received.getAckID() != null) {
+                        receivedMessages.add(received);
+                        // Ack immediately after receiving (not in batch)
+                        subscription.sendAck(received.getAckID());
+                        System.out.println("Received and acked message " + receivedMessages.size() + "/" + numMessages);
+                    } else {
+                        TimeUnit.MILLISECONDS.sleep(100);
+                    }
+                } catch (Exception e) {
+                    System.err.println("Error receiving message: " + e.getMessage());
+                    TimeUnit.MILLISECONDS.sleep(100);
+                }
+            }
+
+            Assertions.assertEquals(numMessages, receivedMessages.size(),
+                    "Should receive all messages. Expected: " + numMessages + ", Got: " + receivedMessages.size());
+
+            // Verify all messages were received
+            for (int i = 0; i < receivedMessages.size(); i++) {
+                Message received = receivedMessages.get(i);
+                Assertions.assertNotNull(received, "Received message " + i + " should not be null");
+                Assertions.assertNotNull(received.getBody(), "Received message " + i + " body should not be null");
+                Assertions.assertNotNull(received.getAckID(), "Received message " + i + " should have AckID");
+            }
+        }
+    }
 }

pubsub/pubsub-client/src/test/java/com/salesforce/multicloudj/pubsub/driver/AbstractSubscriptionTest.java

@@ -18,6 +18,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -713,4 +714,47 @@ void testBuilderProviderId() {
 
         assertEquals("my-provider", sub.getProviderId());
     }
+
+    @Test
+    void testBatchSizeNeverZero() throws Exception {
+        // Test that updateBatchSize() never returns 0
+        MockMessageSource source = new MockMessageSource();
+        List<Integer> batchSizes = Collections.synchronizedList(new ArrayList<>());
+        final AtomicInteger callCount = new AtomicInteger(0);
+        
+        TestSubscription sub = new TestSubscription(source) {
+            @Override
+            protected List<Message> doReceiveBatch(int batchSize) {
+                batchSizes.add(batchSize);
+                int call = callCount.incrementAndGet();
+                
+                // First call: fetch the message so receive() doesn't hang
+                // Subsequent calls: return empty to trigger batch size shrinking
+                if (call == 1) {
+                    return source.fetchBatch(batchSize);
+                }
+                return Collections.emptyList();
+            }
+        };
+
+        // Add one message to trigger prefetch
+        source.addMessages(List.of(
+            Message.builder().withBody("test".getBytes()).build()
+        ));
+
+        // Receive the message - this triggers prefetch which calls updateBatchSize()
+        Message msg = sub.receive();
+        assertNotNull(msg);
+
+        // Wait for prefetch operations to complete
+        Thread.sleep(500);
+
+        // All batch sizes should be >= 1
+        assertFalse(batchSizes.isEmpty(), "doReceiveBatch should have been called at least once");
+        for (int size : batchSizes) {
+            assertTrue(size >= 1, "Batch size from updateBatchSize() should never be less than 1, but got: " + size);
+        }
+
+        sub.close();
+    }
 }