Prevent improper error logging during worker shutdown (awslabs#1257)

* Move throwOnIllegalState call to drain queue method to prevent improper error logging during worker shutdown

* Fix unit tests that expected IllegalStateException thrown

* Changed names of unit tests to reflect new behavior

amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/polling/PrefetchRecordsPublisher.java

@@ -289,7 +289,6 @@ private void throwOnIllegalState() {
     }
 
     private PrefetchRecordsRetrieved peekNextResult() {
-        throwOnIllegalState();
         return publisherSession.peekNextRecord();
     }
 
@@ -336,6 +335,7 @@ public void restartFrom(RecordsRetrieved recordsRetrieved) {
 
     @Override
     public void subscribe(Subscriber<? super RecordsRetrieved> s) {
+        throwOnIllegalState();
         subscriber = s;
         subscriber.onSubscribe(new Subscription() {
             @Override
@@ -389,6 +389,7 @@ synchronized void drainQueueForRequests() {
         // If there is an event available to drain and if there is at least one demand,
         // then schedule it for delivery
         if (publisherSession.hasDemandToPublish() && canDispatchRecord(recordsToDeliver)) {
+            throwOnIllegalState();
             subscriber.onNext(recordsToDeliver.prepareForPublish());
             recordsToDeliver.dispatched();
             lastEventDeliveryTime = Instant.now();

amazon-kinesis-client/src/test/java/software/amazon/kinesis/retrieval/polling/PrefetchRecordsPublisherTest.java

@@ -31,8 +31,10 @@
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -64,6 +66,7 @@
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.runners.MockitoJUnitRunner;
@@ -375,15 +378,29 @@ record = Record.builder().data(createByteBufferWithSize(1024)).build();
     }
 
     @Test(expected = IllegalStateException.class)
-    public void testGetNextRecordsWithoutStarting() {
+    public void testSubscribeWithoutStarting() {
         verify(executorService, never()).execute(any());
-        getRecordsCache.drainQueueForRequests();
+        Subscriber<RecordsRetrieved> mockSubscriber = mock(Subscriber.class);
+        getRecordsCache.subscribe(mockSubscriber);
     }
 
     @Test(expected = IllegalStateException.class)
-    public void testCallAfterShutdown() {
+    public void testRequestRecordsOnSubscriptionAfterShutdown() {
+        GetRecordsResponse response = GetRecordsResponse.builder().records(
+                        Record.builder().data(SdkBytes.fromByteArray(new byte[] { 1, 2, 3 })).sequenceNumber("123").build())
+                .nextShardIterator(NEXT_SHARD_ITERATOR).build();
+        when(getRecordsRetrievalStrategy.getRecords(anyInt())).thenReturn(response);
+
+        getRecordsCache.start(sequenceNumber, initialPosition);
+
+        verify(getRecordsRetrievalStrategy, timeout(100).atLeastOnce()).getRecords(anyInt());
+
         when(executorService.isShutdown()).thenReturn(true);
-        getRecordsCache.drainQueueForRequests();
+        Subscriber<RecordsRetrieved> mockSubscriber = mock(Subscriber.class);
+        getRecordsCache.subscribe(mockSubscriber);
+        ArgumentCaptor<Subscription> subscriptionCaptor = ArgumentCaptor.forClass(Subscription.class);
+        verify(mockSubscriber).onSubscribe(subscriptionCaptor.capture());
+        subscriptionCaptor.getValue().request(1);
     }
 
     @Test