Added a listener for the connection state to start writing at the correct time (#1121)

Author: codepitbull

hivemq-edge/src/main/java/com/hivemq/edge/modules/adapters/impl/ProtocolAdapterStateImpl.java

@@ -25,6 +25,7 @@
 import org.apache.commons.lang3.exception.ExceptionUtils;
 
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 public class ProtocolAdapterStateImpl implements ProtocolAdapterState {
     private final @NotNull EventService eventService;
@@ -34,6 +35,7 @@ public class ProtocolAdapterStateImpl implements ProtocolAdapterState {
     protected @NotNull AtomicReference<ConnectionStatus> connectionStatus =
             new AtomicReference<>(ConnectionStatus.DISCONNECTED);
     protected @Nullable String lastErrorMessage;
+    private AtomicReference<Consumer<ConnectionStatus>> connectionStatusListener = new AtomicReference<>();
 
     public ProtocolAdapterStateImpl(final @NotNull EventService eventService,
                                     final @NotNull String adapterId,
@@ -46,7 +48,14 @@ public ProtocolAdapterStateImpl(final @NotNull EventService eventService,
     @Override
     public boolean setConnectionStatus(final @NotNull ConnectionStatus connectionStatus) {
         Preconditions.checkNotNull(connectionStatus);
-        return this.connectionStatus.getAndSet(connectionStatus) != connectionStatus;
+        final var changed = this.connectionStatus.getAndSet(connectionStatus) != connectionStatus;
+        if(changed) {
+            final var listener = connectionStatusListener.get();
+            if(listener != null) {
+                listener.accept(connectionStatus);
+            }
+        }
+        return changed;
     }
 
     @Override
@@ -102,4 +111,9 @@ public void setRuntimeStatus(final @NotNull RuntimeStatus runtimeStatus) {
         return lastErrorMessage;
     }
 
+    public void setConnectionStatusListener(Consumer<ConnectionStatus> connectionStatusListener) {
+        this.connectionStatusListener.set(connectionStatusListener);
+        connectionStatusListener.accept(connectionStatus.get());
+    }
+
 }

hivemq-edge/src/main/java/com/hivemq/protocols/ProtocolAdapterWrapper.java

@@ -31,6 +31,7 @@
 import com.hivemq.adapter.sdk.api.tag.Tag;
 import com.hivemq.adapter.sdk.api.writing.WritingProtocolAdapter;
 import com.hivemq.edge.modules.adapters.data.TagManager;
+import com.hivemq.edge.modules.adapters.impl.ProtocolAdapterStateImpl;
 import com.hivemq.edge.modules.api.adapters.ProtocolAdapterPollingService;
 import com.hivemq.persistence.mappings.NorthboundMapping;
 import com.hivemq.persistence.mappings.SouthboundMapping;
@@ -48,6 +49,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -69,7 +71,7 @@ private enum OperationState {
     private final @NotNull ProtocolAdapter adapter;
     private final @NotNull ProtocolAdapterFactory<?> adapterFactory;
     private final @NotNull ProtocolAdapterInformation adapterInformation;
-    private final @NotNull ProtocolAdapterState protocolAdapterState;
+    private final @NotNull ProtocolAdapterStateImpl protocolAdapterState;
     private final @NotNull InternalProtocolAdapterWritingService protocolAdapterWritingService;
     private final @NotNull ProtocolAdapterPollingService protocolAdapterPollingService;
     private final @NotNull ProtocolAdapterConfig config;
@@ -109,7 +111,7 @@ public ProtocolAdapterWrapper(
             final @NotNull ProtocolAdapter adapter,
             final @NotNull ProtocolAdapterFactory<?> adapterFactory,
             final @NotNull ProtocolAdapterInformation adapterInformation,
-            final @NotNull ProtocolAdapterState protocolAdapterState,
+            final @NotNull ProtocolAdapterStateImpl protocolAdapterState,
             final @NotNull NorthboundConsumerFactory northboundConsumerFactory,
             final @NotNull TagManager tagManager) {
         this.protocolAdapterWritingService = protocolAdapterWritingService;
@@ -202,12 +204,20 @@ private Optional<Throwable> attemptStartingConsumers(final boolean writingEnable
             //Adapter started successfully, now start the consumers
             createAndSubscribeTagConsumer();
             startPolling(protocolAdapterPollingService, eventService);
-            if(startWriting(writingEnabled, protocolAdapterWritingService)) {
-                log.info("Successfully started adapter with id {}", adapter.getId());
-            } else {
-                log.error("Protocol adapter start failed as data hub is not available.");
-                return Optional.of(new RuntimeException(
-                        "Protocol adapter start failed as data hub is not available."));
+
+            if(writingEnabled && isWriting()) {
+                final var started = new AtomicBoolean(false);
+                protocolAdapterState.setConnectionStatusListener(status -> {
+                    if(status == ProtocolAdapterState.ConnectionStatus.CONNECTED) {
+                        if(started.compareAndSet(false, true)) {
+                            if(startWriting(protocolAdapterWritingService)) {
+                                log.info("Successfully started adapter with id {}", adapter.getId());
+                            } else {
+                                log.error("Protocol adapter start failed as data hub is not available.");
+                            }
+                        }
+                    }
+                });
             }
         } catch (final Throwable e) {
             log.error("Protocol adapter start failed");
@@ -432,25 +442,19 @@ private void stopPolling(
         }
     }
 
-    private @NotNull boolean startWriting(
-            final boolean writingEnabled,
-            final @NotNull InternalProtocolAdapterWritingService protocolAdapterWritingService) {
-        if (writingEnabled && isWriting()) {
-            log.debug("Start writing for protocol adapter with id '{}'", getId());
-
-            final var southboundMappings = getSouthboundMappings();
-            final var writingContexts = southboundMappings.stream()
-                            .map(InternalWritingContextImpl::new)
-                            .collect(Collectors.<InternalWritingContext>toList());
-
-            return protocolAdapterWritingService
-                    .startWriting(
-                        (WritingProtocolAdapter) getAdapter(),
-                        getProtocolAdapterMetricsService(),
-                        writingContexts);
-        } else {
-            return true;
-        }
+    private @NotNull boolean startWriting(final @NotNull InternalProtocolAdapterWritingService protocolAdapterWritingService) {
+        log.debug("Start writing for protocol adapter with id '{}'", getId());
+
+        final var southboundMappings = getSouthboundMappings();
+        final var writingContexts = southboundMappings.stream()
+                        .map(InternalWritingContextImpl::new)
+                        .collect(Collectors.<InternalWritingContext>toList());
+
+        return protocolAdapterWritingService
+                .startWriting(
+                    (WritingProtocolAdapter) getAdapter(),
+                    getProtocolAdapterMetricsService(),
+                    writingContexts);
     }
 
     private void stopWriting(final @NotNull InternalProtocolAdapterWritingService protocolAdapterWritingService) {