[#112] Give each Jetty connection its own endpoint and close cached principals

OpenICF-java-framework/connector-server-jetty/src/main/java/org/forgerock/openicf/framework/server/jetty/OpenICFWebSocket.java

@@ -0,0 +1,229 @@
+/*
+ * The contents of this file are subject to the terms of the Common Development and
+ * Distribution License (the License). You may not use this file except in compliance with the
+ * License.
+ *
+ * You can obtain a copy of the License at legal/CDDLv1.0.txt. See the License for the
+ * specific language governing permission and limitations under the License.
+ *
+ * When distributing Covered Software, include this CDDL Header Notice in each file and include
+ * the License file at legal/CDDLv1.0.txt. If applicable, add the following below the CDDL
+ * Header, with the fields enclosed by brackets [] replaced by your own identifying
+ * information: "Portions copyright [year] [name of copyright owner]".
+ *
+ * Copyright 2015-2016 ForgeRock AS.
+ * Portions Copyrighted 2026 3A Systems, LLC
+ */
+package org.forgerock.openicf.framework.server.jetty;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.eclipse.jetty.websocket.api.BatchMode;
+import org.eclipse.jetty.websocket.api.Frame;
+import org.eclipse.jetty.websocket.api.Session;
+import org.eclipse.jetty.websocket.api.StatusCode;
+import org.eclipse.jetty.websocket.api.WebSocketFrameListener;
+import org.eclipse.jetty.websocket.api.WebSocketListener;
+import org.eclipse.jetty.websocket.api.WebSocketPingPongListener;
+import org.forgerock.openicf.common.protobuf.RPCMessages;
+import org.forgerock.openicf.framework.remote.ConnectionPrincipal;
+import org.forgerock.openicf.framework.remote.rpc.RemoteOperationContext;
+import org.forgerock.openicf.framework.remote.rpc.WebSocketConnectionHolder;
+import org.forgerock.util.Utils;
+import org.forgerock.util.promise.Promises;
+import org.identityconnectors.common.logging.Log;
+import org.identityconnectors.framework.common.exceptions.ConnectorIOException;
+
+/**
+ * The Jetty endpoint of a single WebSocket connection.
+ * <p>
+ * One instance is created per connection by {@link OpenICFWebSocketCreator},
+ * while the {@link ConnectionPrincipal} it belongs to is cached per principal
+ * name and shared by every connection of that name. Keeping the session, the
+ * {@link WebSocketConnectionHolder} and the close state here (instead of on
+ * the shared principal) lets concurrent and successive connections of one
+ * principal coexist: each connection delivers its own close event and cleans
+ * up its own resources.
+ */
+public class OpenICFWebSocket implements
+        WebSocketPingPongListener, WebSocketListener, WebSocketFrameListener {
+
+    private static final Log logger = Log.getLog(OpenICFWebSocket.class);
+
+    private final ConnectionPrincipal<?> principal;
+
+    // Written on the Jetty connect thread, read on the send-executor thread
+    // and by isOperational() callers.
+    private volatile Session session;
+
+    private final AtomicBoolean closed = new AtomicBoolean(false);
+
+    // Written on the handshake-processing pool thread, read by other message
+    // threads via getRemoteConnectionContext()/isHandHooked().
+    private volatile RemoteOperationContext context = null;
+
+    // Single send thread per connection: frames must leave in submission
+    // order (the peer drops e.g. an operation response that overtakes the
+    // handshake response). Jetty's RemoteEndpoint is thread-safe, but
+    // concurrent blocking sends may reach the wire in any order. The executor
+    // is shut down when this connection closes.
+    private final ExecutorService sendExecutor = Executors.newSingleThreadExecutor(
+            Utils.newThreadFactory(null, "OpenICF Jetty WebSocket Send %d", true));
+
+    public OpenICFWebSocket(final ConnectionPrincipal<?> principal) {
+        this.principal = principal;
+    }
+
+    Session getSession() {
+        return session;
+    }
+
+    @Override
+    public void onWebSocketConnect(Session session) {
+        WebSocketPingPongListener.super.onWebSocketConnect(session);
+        this.session = session;
+        principal.getOperationMessageListener().onConnect(adapter);
+    }
+
+    @Override
+    public void onWebSocketClose(int statusCode, String reason) {
+        if (!closed.compareAndSet(false, true)) {
+            return;
+        }
+        try {
+            principal.getOperationMessageListener().onClose(adapter, statusCode, reason);
+        } finally {
+            // Notifies the holder's close listeners so the
+            // WebSocketConnectionGroup drops this connection.
+            adapter.close();
+            sendExecutor.shutdown();
+        }
+    }
+
+    @Override
+    public void onWebSocketError(Throwable t) {
+        logger.warn(t, "onError");
+        principal.getOperationMessageListener().onError(t);
+    }
+
+    @Override
+    public void onWebSocketPing(ByteBuffer buffer) {
+        byte[] b = new byte[buffer.remaining()];
+        buffer.get(b);
+        principal.getOperationMessageListener().onPing(adapter, b);
+    }
+
+    @Override
+    public void onWebSocketPong(ByteBuffer buffer) {
+        byte[] b = new byte[buffer.remaining()];
+        buffer.get(b);
+        principal.getOperationMessageListener().onPong(adapter, b);
+    }
+
+    @Override
+    public void onWebSocketBinary(byte[] payload, int offset, int len) {
+        logger.ok("onBinaryMessage({0})", null != payload ? payload.length : 0);
+        principal.getOperationMessageListener().onMessage(adapter, payload);
+    }
+
+    @Override
+    public void onWebSocketText(String message) {
+        logger.ok("onTextMessage({0})", message);
+        principal.getOperationMessageListener().onMessage(adapter, message);
+    }
+
+    @Override
+    public void onWebSocketFrame(Frame frame) {
+        logger.ok("onWebSocketFrame({0})", frame);
+    }
+
+    private final WebSocketConnectionHolder adapter = new WebSocketConnectionHolder() {
+
+        @Override
+        protected void handshake(RPCMessages.HandshakeMessage message) {
+            context = principal.handshake(this, message);
+        }
+
+        @Override
+        public boolean isOperational() {
+            return null != getSession() && getSession().isOpen();
+        }
+
+        @Override
+        public RemoteOperationContext getRemoteConnectionContext() {
+            return context;
+        }
+
+        @Override
+        public Future<?> sendBytes(byte[] data) {
+            return submitSend(() -> getSession().getRemote().sendBytes(ByteBuffer.wrap(data)));
+        }
+
+        @Override
+        public Future<?> sendString(String data) {
+            return submitSend(() -> getSession().getRemote().sendString(data));
+        }
+
+        private Future<?> submitSend(FrameWrite frame) {
+            if (isOperational()) {
+                try {
+                    return sendExecutor.submit(() -> {
+                        try {
+                            frame.write();
+                        } catch (IOException e) {
+                            throw new RuntimeException(e);
+                        }
+                    });
+                } catch (RejectedExecutionException e) {
+                    // The connection was closed and the executor shut down.
+                }
+            }
+            return Promises.newExceptionPromise(new ConnectorIOException(
+                    "Socket is not connected."));
+        }
+
+        @Override
+        public void sendPing(byte[] applicationData) throws Exception {
+            if (isOperational()) {
+                getSession().getRemote().sendPing(ByteBuffer.wrap(applicationData));
+                if (getSession().getRemote().getBatchMode() == BatchMode.ON) {
+                    getSession().getRemote().flush();
+                }
+            } else {
+                throw new ConnectorIOException("Socket is not connected.");
+            }
+        }
+
+        @Override
+        public void sendPong(byte[] applicationData) throws Exception {
+            if (isOperational()) {
+                getSession().getRemote().sendPong(ByteBuffer.wrap(applicationData));
+                if (getSession().getRemote().getBatchMode() == BatchMode.ON) {
+                    getSession().getRemote().flush();
+                }
+            } else {
+                throw new ConnectorIOException("Socket is not connected.");
+            }
+        }
+
+        @Override
+        protected void tryClose() {
+            final Session current = getSession();
+            if (null != current && current.isOpen()) {
+                current.close(StatusCode.NORMAL, "Shutdown");
+            }
+        }
+
+    };
+
+    @FunctionalInterface
+    private interface FrameWrite {
+        void write() throws IOException;
+    }
+}

OpenICF-java-framework/connector-server-jetty/src/main/java/org/forgerock/openicf/framework/server/jetty/OpenICFWebSocketCreator.java

@@ -12,15 +12,17 @@
  * information: "Portions copyright [year] [name of copyright owner]".
  *
  * Copyright 2015-2016 ForgeRock AS.
- * Portions copyright 2025 3A Systems LLC.
+ * Portions copyright 2025-2026 3A Systems LLC.
  */
 
 package org.forgerock.openicf.framework.server.jetty;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 
 import javax.security.auth.callback.NameCallback;
@@ -39,7 +41,7 @@
 import org.forgerock.openicf.framework.remote.rpc.OperationMessageListener;
 import org.forgerock.openicf.framework.remote.rpc.WebSocketConnectionGroup;
 
-public class OpenICFWebSocketCreator implements JettyWebSocketCreator {
+public class OpenICFWebSocketCreator implements JettyWebSocketCreator, Closeable {
 
     private static final Logger logger = Log.getLogger(OpenICFWebSocketCreator.class);
 
@@ -52,6 +54,12 @@ public class OpenICFWebSocketCreator implements JettyWebSocketCreator {
 
     private Authenticator authenticator;
 
+    private final ScheduledFuture<?> groupHealthChecker;
+
+    // Set by close(); createWebSocket()/authenticate() must not mint new
+    // principals over a framework that is being released.
+    private volatile boolean closed = false;
+
 
     public OpenICFWebSocketCreator(final ConnectorFramework connectorFramework,
                                    final ScheduledExecutorService executorService) {
@@ -84,8 +92,8 @@ public void authenticate(JettyServerUpgradeRequest request, JettyServerUpgradeRe
             this.authenticator = authenticator;
         }
 
-        //Will be cancelled when executorService is shut down
-        executorService.scheduleWithFixedDelay(new Runnable() {
+        //Cancelled in close(); also dies when executorService is shut down
+        groupHealthChecker = executorService.scheduleWithFixedDelay(new Runnable() {
             public void run() {
                 for (WebSocketConnectionGroup e : globalConnectionGroups.values()) {
                     if (!e.checkIsActive()) {
@@ -99,19 +107,61 @@ public void run() {
     @Override
     public Object createWebSocket(JettyServerUpgradeRequest request, JettyServerUpgradeResponse response) {
 
+        if (closed) {
+            if (!response.isCommitted()) {
+                try {
+                    response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE,
+                            "OpenICF connector server is shut down");
+                } catch (IOException e) {
+                    //
+                }
+            }
+            return null;
+        }
+
         if (request.getSubProtocols().contains(RemoteWSFrameworkConnectionInfo.OPENICF_PROTOCOL)) {
             response.setAcceptedSubProtocol(RemoteWSFrameworkConnectionInfo.OPENICF_PROTOCOL);
         }
 
         ConnectionPrincipal<?> connectionPrincipal = authenticate(request, response);
         if (null != connectionPrincipal) {
-            return connectionPrincipal;
+            // The principal is cached and shared by every connection of this
+            // name; the endpoint holds the per-connection state.
+            return new OpenICFWebSocket(connectionPrincipal);
         } else if (!response.isCommitted()) {
             unauthorized(response, "Unknown Principal");
         }
         return null;
     }
 
+    /**
+     * Ends the lifecycle of every cached principal: shuts their connection
+     * groups down and notifies the principals' close listeners. Invoked from
+     * {@link OpenICFWebSocketServletBase#destroy()}.
+     */
+    @Override
+    public void close() {
+        closed = true;
+        groupHealthChecker.cancel(false);
+        // One pass over the groups: every cached principal leaves every group
+        // before the group is dropped, so pending requests of all principals
+        // are cancelled (removing the group inside a per-principal loop would
+        // hide it from the remaining principals).
+        for (WebSocketConnectionGroup group : globalConnectionGroups.values()) {
+            for (ConnectionPrincipal<?> principal : principalCache.values()) {
+                // We should gracefully shut down the group
+                group.principalIsShuttingDown(principal);
+            }
+            if (!group.isOperational()) {
+                globalConnectionGroups.remove(group.getRemoteSessionId());
+            }
+        }
+        for (ConnectionPrincipal<?> principal : principalCache.values()) {
+            principal.close();
+        }
+        principalCache.clear();
+    }
+
     protected void unauthorized(JettyServerUpgradeResponse response, String message) {
         try {
             response.sendError(
@@ -123,17 +173,15 @@ protected void unauthorized(JettyServerUpgradeResponse response, String message)
     }
 
     public ConnectionPrincipal<?> authenticate(JettyServerUpgradeRequest request, JettyServerUpgradeResponse response) {
+        if (closed) {
+            return null;
+        }
         NameCallback callback = new NameCallback("OpenICF user:>");
         authenticator.authenticate(request, response, callback);
         if (StringUtil.isNotBlank(callback.getName())) {
-            ConnectionPrincipal<?> connectionPrincipal = principalCache.get(callback.getName());
-            if (connectionPrincipal == null) {
-                principalCache.putIfAbsent(callback.getName(), new SinglePrincipal(callback.getName(), listener,
-                        connectorFramework, globalConnectionGroups));
-                return principalCache.get(callback.getName());
-            } else {
-                return connectionPrincipal;
-            }
+            return principalCache.computeIfAbsent(callback.getName(),
+                    name -> new SinglePrincipal(name, listener, connectorFramework,
+                            globalConnectionGroups));
         }
         return null;
     }