Fix beacon shell session cache concurrency and reuse

app/src/main/java/io/xpipe/app/beacon/AppBeaconCache.java

@@ -6,46 +6,71 @@
 
 import lombok.Value;
 
-import java.util.HashSet;
-import java.util.Set;
+import java.util.Collection;
+import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 
 @Value
 public class AppBeaconCache {
 
-    Set<BeaconShellSession> shellSessions = new HashSet<>();
+    Map<UUID, BeaconShellSession> shellSessionsById = new ConcurrentHashMap<>();
+
+    public Collection<BeaconShellSession> getShellSessions() {
+        return shellSessionsById.values();
+    }
 
     public BeaconShellSession getShellSession(UUID uuid) throws BeaconClientException {
-        var found = shellSessions.stream()
-                .filter(beaconShellSession ->
-                        beaconShellSession.getEntry().getUuid().equals(uuid))
-                .findFirst();
-        if (found.isEmpty()) {
+        var session = shellSessionsById.get(uuid);
+        if (session == null) {
             throw new BeaconClientException("No active shell session known for id " + uuid);
         }
-        return found.get();
+        return session;
     }
 
     public BeaconShellSession getOrStart(DataStoreEntryRef<ShellStore> ref) throws Exception {
-        var existing = AppBeaconServer.get().getCache().getShellSessions().stream()
-                .filter(beaconShellSession -> beaconShellSession.getEntry().equals(ref.get()))
-                .findFirst();
-        var control = (existing.isPresent()
-                ? existing.get().getControl()
-                : ref.getStore().standaloneControl().start());
-        control.setNonInteractive();
-        control.start();
-
-        var d = control.getShellDialect().getDumbMode();
-        if (!d.supportsAnyPossibleInteraction()) {
-            control.close();
-            d.throwIfUnsupported();
+        var entry = ref.get();
+        try {
+            var session = shellSessionsById.computeIfAbsent(entry.getUuid(), key -> {
+                try {
+                    var control = ref.getStore().standaloneControl();
+                    control.setNonInteractive();
+                    control = control.start();
+
+                    var d = control.getShellDialect().getDumbMode();
+                    if (!d.supportsAnyPossibleInteraction()) {
+                        control.close();
+                        d.throwIfUnsupported();
+                    }
+
+                    return new BeaconShellSession(entry, control);
+                } catch (Exception ex) {
+                    throw new ShellSessionInitException(ex);
+                }
+            });
+            session.getControl().setNonInteractive();
+            return session;
+        } catch (ShellSessionInitException ex) {
+            throw ex.getCause();
         }
+    }
 
-        if (existing.isEmpty()) {
-            AppBeaconServer.get().getCache().getShellSessions().add(new BeaconShellSession(ref.get(), control));
+    public void removeShellSession(UUID uuid) {
+        shellSessionsById.remove(uuid);
+    }
+
+    private static final class ShellSessionInitException extends RuntimeException {
+
+        private final Exception cause;
+
+        private ShellSessionInitException(Exception cause) {
+            super(cause);
+            this.cause = cause;
         }
 
-        return new BeaconShellSession(ref.get(), control);
+        @Override
+        public synchronized Exception getCause() {
+            return cause;
+        }
     }
 }

app/src/main/java/io/xpipe/app/beacon/impl/ShellStartExchangeImpl.java

@@ -1,9 +1,9 @@
 package io.xpipe.app.beacon.impl;
 
 import io.xpipe.app.beacon.AppBeaconServer;
-import io.xpipe.app.beacon.BeaconShellSession;
 import io.xpipe.app.ext.ShellStore;
 import io.xpipe.app.storage.DataStorage;
+import io.xpipe.app.storage.DataStoreEntryRef;
 import io.xpipe.beacon.BeaconClientException;
 import io.xpipe.beacon.api.ShellStartExchange;
 import io.xpipe.core.JacksonMapper;
@@ -19,28 +19,13 @@ public Object handle(HttpExchange exchange, Request msg) {
         var e = DataStorage.get()
                 .getStoreEntryIfPresent(msg.getConnection())
                 .orElseThrow(() -> new BeaconClientException("Unknown connection"));
-        if (!(e.getStore() instanceof ShellStore s)) {
+        if (!(e.getStore() instanceof ShellStore)) {
             throw new BeaconClientException("Not a shell connection");
         }
 
-        var existing = AppBeaconServer.get().getCache().getShellSessions().stream()
-                .filter(beaconShellSession -> beaconShellSession.getEntry().equals(e))
-                .findFirst();
-        var control = (existing.isPresent()
-                ? existing.get().getControl()
-                : s.standaloneControl().start());
-        control.setNonInteractive();
-        control.start();
-
-        var d = control.getShellDialect().getDumbMode();
-        if (!d.supportsAnyPossibleInteraction()) {
-            control.close();
-            d.throwIfUnsupported();
-        }
-
-        if (existing.isEmpty()) {
-            AppBeaconServer.get().getCache().getShellSessions().add(new BeaconShellSession(e, control));
-        }
+        var shellSession =
+                AppBeaconServer.get().getCache().getOrStart(new DataStoreEntryRef<ShellStore>(e));
+        var control = shellSession.getControl();
         var ttyState = JacksonMapper.getDefault().valueToTree(control.getTtyState()).asText();
         return Response.builder()
                 .shellDialect(control.getShellDialect().getId())

app/src/main/java/io/xpipe/app/beacon/impl/ShellStopExchangeImpl.java

@@ -13,7 +13,7 @@ public class ShellStopExchangeImpl extends ShellStopExchange {
     public Object handle(HttpExchange exchange, Request msg) {
         var e = AppBeaconServer.get().getCache().getShellSession(msg.getConnection());
         e.getControl().close();
-        AppBeaconServer.get().getCache().getShellSessions().remove(e);
+        AppBeaconServer.get().getCache().removeShellSession(msg.getConnection());
         return Response.builder().build();
     }
 }