From 1c434558e73ccd91112ae1f861eacfe1c1be6db0 Mon Sep 17 00:00:00 2001
From: Enrico Olivelli <enrico.olivelli@datastax.com>
Date: Fri, 26 Jan 2024 12:06:36 +0100
Subject: [PATCH] docker run: prevent concurrent unpacking of nar files

---
 .../java/ai/langstream/impl/nar/NarFileHandler.java  | 12 ++++++++++++
 .../ai/langstream/runtime/agent/AgentRunner.java     |  9 ++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/langstream-core/src/main/java/ai/langstream/impl/nar/NarFileHandler.java b/langstream-core/src/main/java/ai/langstream/impl/nar/NarFileHandler.java
index afd41e8a4..a392548b9 100644
--- a/langstream-core/src/main/java/ai/langstream/impl/nar/NarFileHandler.java
+++ b/langstream-core/src/main/java/ai/langstream/impl/nar/NarFileHandler.java
@@ -210,6 +210,18 @@ public synchronized void scan() throws Exception {
 
     public void handleNarFile(Path narFile) throws Exception {
         String filename = narFile.getFileName().toString();
+        if (packages.containsKey(filename)) {
+            log.error(
+                    "NarFileHandler ID: {} The package {} has already been processed",
+                    System.identityHashCode(this),
+                    filename);
+            log.error(
+                    "NarFileHandler ID: {} Current packages: {}",
+                    System.identityHashCode(this),
+                    packages.keySet());
+            throw new IllegalStateException(
+                    "The package " + filename + " has already been processed");
+        }
 
         // first of all we look for an index file
         try (ZipFile zipFile = new ZipFile(narFile.toFile())) {
diff --git a/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java b/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java
index 509502537..69e5dd267 100644
--- a/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java
+++ b/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java
@@ -160,16 +160,19 @@ public void run(
         log.info("Code directory {}", codeDirectory);
         log.info("Base persistent state directory {}", basePersistentStateDirectory);
 
+        boolean sharingNarFileHandler = sharedNarFileHandler != null;
         List<URL> customLibClasspath = buildCustomLibClasspath(codeDirectory);
         NarFileHandler narFileHandler =
-                sharedNarFileHandler != null
+                sharingNarFileHandler
                         ? sharedNarFileHandler
                         : new NarFileHandler(
                                 agentsDirectory,
                                 customLibClasspath,
                                 Thread.currentThread().getContextClassLoader());
         try {
-            narFileHandler.scan();
+            if (!sharingNarFileHandler) {
+                narFileHandler.scan();
+            }
 
             TopicConnectionsRuntimeRegistry topicConnectionsRuntimeRegistry =
                     new TopicConnectionsRuntimeRegistry();
@@ -208,7 +211,7 @@ public void run(
                 topicConnectionsRuntimeWithClassloader.close();
             }
         } finally {
-            if (sharedNarFileHandler == null) {
+            if (!sharingNarFileHandler) {
                 narFileHandler.close();
             }
         }