improve thread safety of ProtocolAdapterStateImpl

Author: marregui

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

@@ -20,38 +20,42 @@
 import com.hivemq.adapter.sdk.api.events.model.Payload;
 import com.hivemq.adapter.sdk.api.state.ProtocolAdapterState;
 import com.hivemq.edge.modules.api.events.model.EventImpl;
+import org.apache.commons.lang3.exception.ExceptionUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import org.apache.commons.lang3.exception.ExceptionUtils;
 
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
 public class ProtocolAdapterStateImpl implements ProtocolAdapterState {
+    protected final @NotNull AtomicReference<RuntimeStatus> runtimeStatus;
+    protected final @NotNull AtomicReference<ConnectionStatus> connectionStatus;
+    protected final @NotNull AtomicReference<@Nullable String> lastErrorMessage;
     private final @NotNull EventService eventService;
     private final @NotNull String adapterId;
     private final @NotNull String protocolId;
-    protected @NotNull AtomicReference<RuntimeStatus> runtimeStatus = new AtomicReference<>(RuntimeStatus.STOPPED);
-    protected @NotNull AtomicReference<ConnectionStatus> connectionStatus =
-            new AtomicReference<>(ConnectionStatus.DISCONNECTED);
-    protected @Nullable String lastErrorMessage;
-    private final @NotNull AtomicReference<Consumer<ConnectionStatus>> connectionStatusListener = new AtomicReference<>();
+    private final @NotNull AtomicReference<Consumer<ConnectionStatus>> connectionStatusListener;
 
-    public ProtocolAdapterStateImpl(final @NotNull EventService eventService,
-                                    final @NotNull String adapterId,
-                                    final @NotNull String protocolId) {
+    public ProtocolAdapterStateImpl(
+            final @NotNull EventService eventService,
+            final @NotNull String adapterId,
+            final @NotNull String protocolId) {
         this.eventService = eventService;
         this.adapterId = adapterId;
         this.protocolId = protocolId;
+        this.runtimeStatus = new AtomicReference<>(RuntimeStatus.STOPPED);
+        this.connectionStatus = new AtomicReference<>(ConnectionStatus.DISCONNECTED);
+        this.lastErrorMessage = new AtomicReference<>(null);
+        this.connectionStatusListener = new AtomicReference<>();
     }
 
     @Override
     public boolean setConnectionStatus(final @NotNull ConnectionStatus connectionStatus) {
         Preconditions.checkNotNull(connectionStatus);
         final var changed = this.connectionStatus.getAndSet(connectionStatus) != connectionStatus;
-        if(changed) {
+        if (changed) {
             final var listener = connectionStatusListener.get();
-            if(listener != null) {
+            if (listener != null) {
                 listener.accept(connectionStatus);
             }
         }
@@ -68,11 +72,8 @@ public boolean setConnectionStatus(final @NotNull ConnectionStatus connectionSta
      * and the errorMessage to that supplied.
      */
     @Override
-    public void setErrorConnectionStatus(
-            final @Nullable Throwable t,
-            final @Nullable String errorMessage) {
-        final boolean changed = setConnectionStatus(ConnectionStatus.ERROR);
-        reportErrorMessage( t, errorMessage, changed);
+    public void setErrorConnectionStatus(final @Nullable Throwable t, final @Nullable String errorMessage) {
+        reportErrorMessage(t, errorMessage, setConnectionStatus(ConnectionStatus.ERROR));
     }
 
     /**
@@ -86,33 +87,40 @@ public void reportErrorMessage(
             final @Nullable Throwable throwable,
             final @Nullable String errorMessage,
             final boolean sendEvent) {
-        this.lastErrorMessage = errorMessage == null ? throwable == null ? null : throwable.getMessage() : errorMessage;
+        final String msg = errorMessage == null ? throwable == null ? null : throwable.getMessage() : errorMessage;
+        this.lastErrorMessage.set(msg);
         if (sendEvent) {
-            eventService.createAdapterEvent(adapterId, protocolId)
+            final var eventBuilder = eventService.createAdapterEvent(adapterId, protocolId)
                     .withSeverity(EventImpl.SEVERITY.ERROR)
-                    .withMessage(String.format("Adapter '%s' encountered an error.", adapterId))
-                    .withPayload(Payload.ContentType.PLAIN_TEXT, ExceptionUtils.getStackTrace(throwable))
-                    .fire();
+                    .withMessage(String.format("Adapter '%s' encountered an error.", adapterId));
+            if (throwable != null) {
+                eventBuilder.withPayload(Payload.ContentType.PLAIN_TEXT, ExceptionUtils.getStackTrace(throwable));
+            } else if (errorMessage != null) {
+                eventBuilder.withPayload(Payload.ContentType.PLAIN_TEXT, errorMessage);
+            }
+            eventBuilder.fire();
         }
     }
 
     @Override
-    public void setRuntimeStatus(final @NotNull RuntimeStatus runtimeStatus) {
-        this.runtimeStatus.set(runtimeStatus);
+    public @NotNull RuntimeStatus getRuntimeStatus() {
+        return this.runtimeStatus.get();
     }
 
     @Override
-    public @NotNull RuntimeStatus getRuntimeStatus() {
-        return this.runtimeStatus.get();
+    public void setRuntimeStatus(final @NotNull RuntimeStatus runtimeStatus) {
+        this.runtimeStatus.set(runtimeStatus);
     }
 
     @Override
     public @Nullable String getLastErrorMessage() {
-        return lastErrorMessage;
+        return lastErrorMessage.get();
     }
 
-    public void setConnectionStatusListener(final @NotNull Consumer<ConnectionStatus> connectionStatusListener) {
-        this.connectionStatusListener.set(connectionStatusListener);
-        connectionStatusListener.accept(connectionStatus.get());
+    public void setConnectionStatusListener(final @NotNull Consumer<ConnectionStatus> listener) {
+        // Capture current status before setting listener to reduce race window
+        final ConnectionStatus currentStatus = connectionStatus.get();
+        connectionStatusListener.set(listener);
+        listener.accept(currentStatus);
     }
 }

hivemq-edge/src/main/java/com/hivemq/fsm/ProtocolAdapterFSM.java

@@ -69,18 +69,21 @@ public enum AdapterStateEnum {
             AdapterStateEnum.STOPPING, Set.of(AdapterStateEnum.STOPPED)
     );
 
-    private final AtomicReference<StateEnum> northboundState = new AtomicReference<>(StateEnum.DISCONNECTED);
-    private final AtomicReference<StateEnum> southboundState = new AtomicReference<>(StateEnum.DISCONNECTED);
-    private final AtomicReference<AdapterStateEnum> adapterState = new AtomicReference<>(AdapterStateEnum.STOPPED);
-
-    private final List<Consumer<State>> stateTransitionListeners = new CopyOnWriteArrayList<>();
+    private final @NotNull AtomicReference<StateEnum> northboundState;
+    private final @NotNull AtomicReference<StateEnum> southboundState;
+    private final @NotNull AtomicReference<AdapterStateEnum> adapterState;
+    private final @NotNull List<Consumer<State>> stateTransitionListeners;
 
     public record State(AdapterStateEnum state, StateEnum northbound, StateEnum southbound) { }
 
-    private final String adapterId;
+    private final @NotNull String adapterId;
 
     public ProtocolAdapterFSM(final @NotNull String adapterId) {
         this.adapterId = adapterId;
+        this.northboundState = new AtomicReference<>(StateEnum.DISCONNECTED);
+        this.southboundState = new AtomicReference<>(StateEnum.DISCONNECTED);
+        this.adapterState = new AtomicReference<>(AdapterStateEnum.STOPPED);
+        this.stateTransitionListeners = new CopyOnWriteArrayList<>();
     }
 
     public abstract boolean onStarting();
@@ -172,7 +175,6 @@ public void accept(final ProtocolAdapterState.ConnectionStatus connectionStatus)
         final var transitionResult = switch (connectionStatus) {
             case CONNECTED ->
                 transitionNorthboundState(StateEnum.CONNECTED)  && startSouthbound();
-
             case CONNECTING -> transitionNorthboundState(StateEnum.CONNECTING);
             case DISCONNECTED -> transitionNorthboundState(StateEnum.DISCONNECTED);
             case ERROR -> transitionNorthboundState(StateEnum.ERROR);
@@ -243,7 +245,7 @@ public void unregisterStateTransitionListener(final @NotNull Consumer<State> sta
         stateTransitionListeners.remove(stateTransitionListener);
     }
 
-    public State currentState() {
+    public @NotNull State currentState() {
         return new State(adapterState.get(), northboundState.get(), southboundState.get());
     }