Explicitly close the FlowReceiver (#31982)

Author: bzablocki

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/BasicAuthJcsmpSessionService.java

@@ -46,6 +46,7 @@ public class BasicAuthJcsmpSessionService extends SessionService {
   private final String password;
   private final String vpnName;
   @Nullable private JCSMPSession jcsmpSession;
+  @Nullable private MessageReceiver messageReceiver;
   private final RetryCallableManager retryCallableManager = RetryCallableManager.create();
 
   /**
@@ -73,21 +74,25 @@ public void connect() {
 
   @Override
   public void close() {
-    if (isClosed()) {
-      return;
-    }
     retryCallableManager.retryCallable(
         () -> {
-          checkStateNotNull(jcsmpSession).closeSession();
+          if (messageReceiver != null) {
+            messageReceiver.close();
+          }
+          if (!isClosed()) {
+            checkStateNotNull(jcsmpSession).closeSession();
+          }
           return 0;
         },
         ImmutableSet.of(IOException.class));
   }
 
   @Override
   public MessageReceiver createReceiver() {
-    return retryCallableManager.retryCallable(
-        this::createFlowReceiver, ImmutableSet.of(JCSMPException.class));
+    this.messageReceiver =
+        retryCallableManager.retryCallable(
+            this::createFlowReceiver, ImmutableSet.of(JCSMPException.class));
+    return this.messageReceiver;
   }
 
   @Override

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/MessageReceiver.java

@@ -49,6 +49,9 @@ public interface MessageReceiver {
    */
   BytesXMLMessage receive() throws IOException;
 
+  /** Closes the message receiver. */
+  void close();
+
   /**
    * Test clients may return {@literal true} to signal that all expected messages have been pulled
    * and the test may complete. Real clients should always return {@literal false}.

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/SolaceMessageReceiver.java

@@ -69,4 +69,11 @@ public BytesXMLMessage receive() throws IOException {
       throw new IOException(e);
     }
   }
+
+  @Override
+  public void close() {
+    if (!isClosed()) {
+      this.flowReceiver.close();
+    }
+  }
 }

sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/MockSessionService.java

@@ -92,6 +92,9 @@ public BytesXMLMessage receive() throws IOException {
       return getRecordFn.apply(counter.getAndIncrement());
     }
 
+    @Override
+    public void close() {}
+
     @Override
     public boolean isEOF() {
       return counter.get() >= minMessagesReceived;