concord-server: refactor WebSocketChannelManager, allow message sources in plugins

Author: ibodrov

server/impl/src/main/java/com/walmartlabs/concord/server/ApiServerModule.java

@@ -28,7 +28,7 @@
 import com.walmartlabs.concord.server.boot.servlets.FormServletHolder;
 import com.walmartlabs.concord.server.boot.statics.StaticResourcesConfigurator;
 import com.walmartlabs.concord.server.boot.validation.ValidationModule;
-import com.walmartlabs.concord.server.websocket.ConcordWebSocketServlet;
+import com.walmartlabs.concord.server.agent.websocket.ConcordWebSocketServlet;
 import org.apache.shiro.mgt.SecurityManager;
 import org.apache.shiro.web.mgt.WebSecurityManager;
 import org.eclipse.jetty.ee8.servlet.FilterHolder;

server/impl/src/main/java/com/walmartlabs/concord/server/ServerResource.java

@@ -25,7 +25,7 @@
 import com.walmartlabs.concord.server.boot.BackgroundTasks;
 import com.walmartlabs.concord.server.sdk.rest.Resource;
 import com.walmartlabs.concord.server.task.TaskScheduler;
-import com.walmartlabs.concord.server.websocket.WebSocketChannelManager;
+import com.walmartlabs.concord.server.agent.message.MessageChannelManager;
 import org.jooq.Configuration;
 
 import javax.inject.Inject;
@@ -42,18 +42,18 @@ public class ServerResource implements Resource {
 
     private final TaskScheduler taskScheduler;
     private final BackgroundTasks backgroundTasks;
-    private final WebSocketChannelManager webSocketChannelManager;
+    private final MessageChannelManager messageChannelManager;
     private final PingDao pingDao;
 
     @Inject
     public ServerResource(TaskScheduler taskScheduler,
                           BackgroundTasks backgroundTasks,
-                          WebSocketChannelManager webSocketChannelManager,
+                          MessageChannelManager messageChannelManager,
                           PingDao pingDao) {
 
         this.taskScheduler = taskScheduler;
         this.backgroundTasks = backgroundTasks;
-        this.webSocketChannelManager = webSocketChannelManager;
+        this.messageChannelManager = messageChannelManager;
         this.pingDao = pingDao;
     }
 
@@ -78,7 +78,7 @@ public VersionResponse version() {
     public void maintenanceMode() {
         backgroundTasks.stop();
 
-        webSocketChannelManager.shutdown();
+        messageChannelManager.shutdown();
         taskScheduler.stop();
     }
 

server/impl/src/main/java/com/walmartlabs/concord/server/agent/AgentManager.java

@@ -20,11 +20,12 @@
  * =====
  */
 
+import com.walmartlabs.concord.server.agent.message.MessageChannel;
+import com.walmartlabs.concord.server.agent.message.MessageChannelManager;
+import com.walmartlabs.concord.server.agent.websocket.WebSocketChannel;
 import com.walmartlabs.concord.server.queueclient.message.MessageType;
 import com.walmartlabs.concord.server.queueclient.message.ProcessRequest;
 import com.walmartlabs.concord.server.sdk.ProcessKey;
-import com.walmartlabs.concord.server.websocket.WebSocketChannel;
-import com.walmartlabs.concord.server.websocket.WebSocketChannelManager;
 import org.jooq.DSLContext;
 
 import javax.inject.Inject;
@@ -38,23 +39,24 @@
 public class AgentManager {
 
     private final AgentCommandsDao commandQueue;
-    private final WebSocketChannelManager channelManager;
+    private final MessageChannelManager channelManager;
 
     @Inject
     public AgentManager(AgentCommandsDao commandQueue,
-                        WebSocketChannelManager channelManager) {
+                        MessageChannelManager channelManager) {
 
         this.commandQueue = commandQueue;
         this.channelManager = channelManager;
     }
 
     public Collection<AgentWorkerEntry> getAvailableAgents() {
-        Map<WebSocketChannel, ProcessRequest> reqs = channelManager.getRequests(MessageType.PROCESS_REQUEST);
+        Map<MessageChannel, ProcessRequest> reqs = channelManager.getRequests(MessageType.PROCESS_REQUEST);
         return reqs.entrySet().stream()
+                .filter(r -> r.getKey() instanceof WebSocketChannel) // TODO a better way
                 .map(r -> AgentWorkerEntry.builder()
                         .channelId(r.getKey().getChannelId())
                         .agentId(r.getKey().getAgentId())
-                        .userAgent(r.getKey().getUserAgent())
+                        .userAgent(((WebSocketChannel) r.getKey()).getUserAgent())
                         .capabilities(r.getValue().getCapabilities())
                         .build())
                 .collect(Collectors.toList());

server/impl/src/main/java/com/walmartlabs/concord/server/agent/AgentWorkerEntry.java

@@ -27,15 +27,14 @@
 
 import javax.annotation.Nullable;
 import java.util.Map;
-import java.util.UUID;
 
 @Value.Immutable
 @JsonInclude(JsonInclude.Include.NON_EMPTY)
 @JsonSerialize(as = ImmutableAgentWorkerEntry.class)
 @JsonDeserialize(as = ImmutableAgentWorkerEntry.class)
 public interface AgentWorkerEntry {
 
-    UUID channelId();
+    String channelId();
 
     @Nullable // backward-compatibility with old agent versions
     String agentId();

server/impl/src/main/java/com/walmartlabs/concord/server/agent/dispatcher/Dispatcher.java

@@ -29,14 +29,14 @@
 import com.walmartlabs.concord.server.PeriodicTask;
 import com.walmartlabs.concord.server.agent.AgentCommand;
 import com.walmartlabs.concord.server.agent.Commands;
+import com.walmartlabs.concord.server.agent.message.MessageChannel;
+import com.walmartlabs.concord.server.agent.message.MessageChannelManager;
 import com.walmartlabs.concord.server.cfg.AgentConfiguration;
 import com.walmartlabs.concord.server.jooq.tables.records.AgentCommandsRecord;
 import com.walmartlabs.concord.server.queueclient.message.CommandRequest;
 import com.walmartlabs.concord.server.queueclient.message.CommandResponse;
 import com.walmartlabs.concord.server.queueclient.message.MessageType;
 import com.walmartlabs.concord.server.sdk.metrics.WithTimer;
-import com.walmartlabs.concord.server.websocket.WebSocketChannel;
-import com.walmartlabs.concord.server.websocket.WebSocketChannelManager;
 import org.jooq.Configuration;
 import org.jooq.DSLContext;
 import org.slf4j.Logger;
@@ -61,12 +61,12 @@ public class Dispatcher extends PeriodicTask {
     private static final int BATCH_SIZE = 10;
 
     private final DispatcherDao dao;
-    private final WebSocketChannelManager channelManager;
+    private final MessageChannelManager channelManager;
 
     @Inject
     public Dispatcher(AgentConfiguration cfg,
                       DispatcherDao dao,
-                      WebSocketChannelManager channelManager) {
+                      MessageChannelManager channelManager) {
 
         super(cfg.getCommandPollDelay().toMillis(), ERROR_DELAY);
         this.dao = dao;
@@ -75,7 +75,7 @@ public Dispatcher(AgentConfiguration cfg,
 
     @Override
     protected boolean performTask() {
-        Map<WebSocketChannel, CommandRequest> requests = this.channelManager.getRequests(MessageType.COMMAND_REQUEST);
+        Map<MessageChannel, CommandRequest> requests = this.channelManager.getRequests(MessageType.COMMAND_REQUEST);
         if (requests.isEmpty()) {
             return false;
         }
@@ -148,7 +148,6 @@ private AgentCommand findCandidate(CommandRequest req, List<AgentCommand> candid
     }
 
     private void sendResponse(Match match, AgentCommand response) {
-        WebSocketChannel channel = match.request.channel;
         long correlationId = match.request.request.getCorrelationId();
 
         CommandType type = CommandType.valueOf((String) response.getData().remove(Commands.TYPE_KEY));
@@ -157,7 +156,8 @@ private void sendResponse(Match match, AgentCommand response) {
         payload.put("type", type.toString());
         payload.putAll(response.getData());
 
-        boolean success = channelManager.sendResponse(channel.getChannelId(), CommandResponse.cancel(correlationId, payload));
+        MessageChannel channel = match.request.channel;
+        boolean success = channelManager.sendMessage(channel.getChannelId(), CommandResponse.cancel(correlationId, payload));
         if (success) {
             log.info("sendResponse ['{}'] -> done", correlationId);
         } else {
@@ -223,25 +223,10 @@ private AgentCommand convert(AgentCommandsRecord r) {
         }
     }
 
-    private static final class Match {
+    private record Match(Request request, AgentCommand command) {
 
-        private final Request request;
-        private final AgentCommand command;
-
-        private Match(Request request, AgentCommand command) {
-            this.request = request;
-            this.command = command;
-        }
     }
 
-    private static final class Request {
-
-        private final WebSocketChannel channel;
-        private final CommandRequest request;
-
-        private Request(WebSocketChannel channel, CommandRequest request) {
-            this.channel = channel;
-            this.request = request;
-        }
+    private record Request(MessageChannel channel, CommandRequest request) {
     }
 }

server/impl/src/main/java/com/walmartlabs/concord/server/agent/message/MessageChannel.java

@@ -0,0 +1,60 @@
+package com.walmartlabs.concord.server.agent.message;
+
+/*-
+ * *****
+ * Concord
+ * -----
+ * Copyright (C) 2017 - 2024 Walmart Inc.
+ * -----
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =====
+ */
+
+
+import com.walmartlabs.concord.server.queueclient.message.Message;
+import com.walmartlabs.concord.server.queueclient.message.MessageType;
+
+import java.io.Closeable;
+import java.util.Optional;
+
+/**
+ * Represents an open message channel between the server and a remote agent.
+ */
+public interface MessageChannel extends Closeable {
+
+    /**
+     * A unique identifier of a channel.
+     */
+    String getChannelId();
+
+    /**
+     * An identifier of the agent represented by this MessageChannel.
+     */
+    String getAgentId();
+
+    /**
+     * Attempts to send a message.
+     * @return {@code true} if the message was sent successfully.
+     * Returns {@code false} if the message cannot be sent at the moment.
+     * @apiNote The implementors must expect the server to re-attempt
+     * the sending of the same message.
+     */
+    boolean offerMessage(Message msg) throws Exception;
+
+    /**
+     * Attempts to grab a received message.
+     * @return an empty value if no messages of the specified type can be returned
+     * at the moment.
+     */
+    Optional<Message> getMessage(MessageType messageType) throws Exception;
+}

server/impl/src/main/java/com/walmartlabs/concord/server/agent/message/MessageChannelManager.java

@@ -0,0 +1,127 @@
+package com.walmartlabs.concord.server.agent.message;
+
+/*-
+ * *****
+ * Concord
+ * -----
+ * Copyright (C) 2017 - 2018 Walmart Inc.
+ * -----
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =====
+ */
+
+import com.walmartlabs.concord.server.queueclient.message.Message;
+import com.walmartlabs.concord.server.queueclient.message.MessageType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.inject.Named;
+import javax.inject.Singleton;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+
+@Named
+@Singleton
+public class MessageChannelManager {
+
+    private static final Logger log = LoggerFactory.getLogger(MessageChannelManager.class);
+
+    private final Map<String, MessageChannel> channels = new ConcurrentHashMap<>();
+
+    private volatile boolean isShutdown;
+
+    public boolean isShutdown() {
+        return isShutdown;
+    }
+
+    public void shutdown() {
+        isShutdown = true;
+
+        channels.forEach((uuid, channel) -> {
+            try {
+                channel.close();
+            } catch (Exception e) {
+                log.warn("shutdown -> failed on channel {}: {}", channel.getClass(), e.getMessage());
+            }
+        });
+        log.info("shutdown -> done");
+    }
+
+    public void close(String channelId) {
+        MessageChannel channel = channels.remove(channelId);
+        if (channel == null) {
+            log.warn("close ['{}'] -> channel not found", channelId);
+            return;
+        }
+
+        try {
+            channel.close();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        log.info("close ['{}'] -> done", channelId);
+    }
+
+    public boolean sendMessage(String channelId, Message response) {
+        MessageChannel channel = channels.get(channelId);
+        if (channel == null) {
+            log.warn("sendResponse ['{}', '{}'] -> channel not found", channelId, response);
+            return false;
+        }
+
+        try {
+            return channel.offerMessage(response);
+        } catch (Exception e) {
+            log.warn("sendResponse ['{}', '{}'] -> failed: {}", channelId, response, e.getMessage());
+            return false;
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    public <M extends Message> Map<MessageChannel, M> getRequests(MessageType requestType) {
+        Map<MessageChannel, M> result = new HashMap<>();
+        channels.forEach((channelId, channel) -> {
+            try {
+                channel.getMessage(requestType).ifPresent(msg -> {
+                    result.put(channel, (M) msg);
+                });
+            } catch (Exception e) {
+                log.warn("getRequests ['{}'] -> failed on channel {}: {}", requestType, channel.getClass(), e.getMessage());
+            }
+        });
+        return result;
+    }
+
+    public void add(MessageChannel channel) {
+        channels.put(channel.getChannelId(), channel);
+    }
+
+    public int connectedClientsCount() {
+        return channels.size();
+    }
+
+    public Map<String, MessageChannel> getChannels() {
+        return Collections.unmodifiableMap(channels);
+    }
+
+    @SuppressWarnings("unchecked")
+    public <C extends MessageChannel> Optional<C> getChannel(String channelId, Class<C> klass) {
+        return Optional.ofNullable(channels.get(channelId))
+                .filter(c -> klass.isAssignableFrom(c.getClass()))
+                .map(c -> (C) c);
+    }
+}