Fix import in multithreaded context

pom.xml

@@ -4,12 +4,12 @@
     <parent>
         <artifactId>pom-scijava</artifactId>
         <groupId>org.scijava</groupId>
-        <version>34.1.0</version>
+        <version>35.1.1</version>
     </parent>
 
     <groupId>fr.igred</groupId>
     <artifactId>simple-omero-client</artifactId>
-    <version>5.12.3</version>
+    <version>5.13.0</version>
     <packaging>jar</packaging>
 
     <name>Simple OMERO Client</name>

src/main/java/fr/igred/omero/GatewayWrapper.java

@@ -26,7 +26,6 @@
 import ome.formats.OMEROMetadataStoreClient;
 import omero.api.IQueryPrx;
 import omero.gateway.Gateway;
-import omero.gateway.JoinSessionCredentials;
 import omero.gateway.LoginCredentials;
 import omero.gateway.SecurityContext;
 import omero.gateway.exception.DSOutOfServiceException;
@@ -43,6 +42,9 @@
 
 import java.util.List;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 
 /**
@@ -52,6 +54,12 @@
  */
 public abstract class GatewayWrapper {
 
+    /** Number of requested import stores */
+    private final AtomicInteger storeUses = new AtomicInteger(0);
+
+    /** Import store lock */
+    private final Lock storeLock = new ReentrantLock(true);
+
     /** Gateway linking the code to OMERO, only linked to one group. */
     private Gateway gateway;
 
@@ -87,6 +95,35 @@ public Gateway getGateway() {
     }
 
 
+    /**
+     * Retrieves the shared import store in a thread-safe way.
+     *
+     * @throws DSOutOfServiceException If the connection is broken, or not logged in.
+     */
+    private OMEROMetadataStoreClient getImportStoreLocked() throws DSOutOfServiceException {
+        storeLock.lock();
+        try {
+            return gateway.getImportStore(ctx);
+        } finally {
+            storeLock.unlock();
+        }
+    }
+
+
+    /**
+     * Closes the import store in a thread-safe manner.
+     */
+    private void closeImportStoreLocked() {
+        if (storeLock.tryLock()) {
+            try {
+                gateway.closeImport(ctx, null);
+            } finally {
+                storeLock.unlock();
+            }
+        }
+    }
+
+
     /**
      * Returns the current user.
      *
@@ -163,7 +200,7 @@ public boolean isConnected() {
      */
     public void connect(String hostname, int port, String sessionId)
     throws ServiceException {
-        connect(new JoinSessionCredentials(sessionId, hostname, port));
+        connect(new LoginCredentials(sessionId, sessionId, hostname, port));
     }
 
 
@@ -233,6 +270,8 @@ public void connect(LoginCredentials cred) throws ServiceException {
     public void disconnect() {
         if (isConnected()) {
             boolean sudo = ctx.isSudo();
+            storeUses.set(0);
+            closeImport();
             user = new ExperimenterWrapper(new ExperimenterData());
             ctx = new SecurityContext(-1);
             ctx.setExperimenter(user.asDataObject());
@@ -348,18 +387,30 @@ public AdminFacility getAdminFacility() throws ExecutionException {
     /**
      * Creates or recycles the import store.
      *
-     * @return config.
+     * @return See above.
      *
      * @throws ServiceException Cannot connect to OMERO.
      */
     public OMEROMetadataStoreClient getImportStore() throws ServiceException {
-        return ExceptionHandler.of(gateway, g -> g.getImportStore(ctx))
+        storeUses.incrementAndGet();
+        return ExceptionHandler.of(this, GatewayWrapper::getImportStoreLocked)
                                .rethrow(DSOutOfServiceException.class, ServiceException::new,
                                         "Could not retrieve import store")
                                .get();
     }
 
 
+    /**
+     * Closes the import store.
+     */
+    public void closeImport() {
+        int remainingStores = storeUses.decrementAndGet();
+        if (remainingStores <= 0) {
+            closeImportStoreLocked();
+        }
+    }
+
+
     /**
      * Finds objects on OMERO through a database query.
      *

src/main/java/fr/igred/omero/repository/DatasetWrapper.java

@@ -463,7 +463,28 @@ public void removeImage(Client client, ImageWrapper image)
      */
     public boolean importImages(Client client, String... paths)
     throws ServiceException, OMEROServerError, AccessException, IOException, ExecutionException {
-        boolean success = importImages(client, data, paths);
+        return importImages(client, 1, paths);
+    }
+
+
+    /**
+     * Imports all images candidates in the paths to the dataset in OMERO.
+     *
+     * @param client  The client handling the connection.
+     * @param threads The number of threads (same value used for filesets and uploads).
+     * @param paths   Paths to the image files on the computer.
+     *
+     * @return If the import did not exit because of an error.
+     *
+     * @throws ServiceException   Cannot connect to OMERO.
+     * @throws AccessException    Cannot access data.
+     * @throws OMEROServerError   Server error.
+     * @throws IOException        Cannot read file.
+     * @throws ExecutionException A Facility can't be retrieved or instantiated.
+     */
+    public boolean importImages(Client client, int threads, String... paths)
+    throws ServiceException, OMEROServerError, AccessException, IOException, ExecutionException {
+        boolean success = importImages(client, data, threads, paths);
         refresh(client);
         return success;
     }

src/main/java/fr/igred/omero/repository/GenericRepositoryObjectWrapper.java

@@ -70,17 +70,18 @@ protected GenericRepositoryObjectWrapper(T o) {
     /**
      * Imports all images candidates in the paths to the target in OMERO.
      *
-     * @param client The client handling the connection.
-     * @param target The import target.
-     * @param paths  Paths to the image files on the computer.
+     * @param client  The client handling the connection.
+     * @param target  The import target.
+     * @param threads The number of threads (same value used for filesets and uploads).
+     * @param paths   Paths to the image files on the computer.
      *
      * @return If the import did not exit because of an error.
      *
      * @throws ServiceException Cannot connect to OMERO.
      * @throws OMEROServerError Server error.
      * @throws IOException      Cannot read file.
      */
-    protected static boolean importImages(GatewayWrapper client, DataObject target, String... paths)
+    protected static boolean importImages(GatewayWrapper client, DataObject target, int threads, String... paths)
     throws ServiceException, OMEROServerError, IOException {
         boolean success;
 
@@ -89,6 +90,8 @@ protected static boolean importImages(GatewayWrapper client, DataObject target,
         config.target.set(type + ":" + target.getId());
         config.username.set(client.getUser().getUserName());
         config.email.set(client.getUser().getEmail());
+        config.parallelFileset.set(threads);
+        config.parallelUpload.set(threads);
 
         OMEROMetadataStoreClient store = client.getImportStore();
         try (OMEROWrapper reader = new OMEROWrapper(config)) {
@@ -107,7 +110,7 @@ protected static boolean importImages(GatewayWrapper client, DataObject target,
             ImportCandidates candidates = new ImportCandidates(reader, paths, handler);
             success = library.importCandidates(config, candidates);
         } finally {
-            store.logout();
+            client.closeImport();
         }
 
         return success;
@@ -163,7 +166,7 @@ protected static List<Long> importImage(GatewayWrapper client, DataObject target
         } catch (Throwable e) {
             throw new OMEROServerError(e);
         } finally {
-            store.logout();
+            client.closeImport();
         }
 
         List<Long> ids = new ArrayList<>(pixels.size());

src/main/java/fr/igred/omero/repository/ScreenWrapper.java

@@ -351,7 +351,28 @@ public void refresh(GatewayWrapper client) throws ServiceException, AccessExcept
      */
     public boolean importImages(GatewayWrapper client, String... paths)
     throws ServiceException, OMEROServerError, AccessException, IOException, ExecutionException {
-        boolean success = importImages(client, data, paths);
+        return this.importImages(client, 1, paths);
+    }
+
+
+    /**
+     * Imports all images candidates in the paths to the screen in OMERO.
+     *
+     * @param client  The client handling the connection.
+     * @param threads The number of threads (same value used for filesets and uploads).
+     * @param paths   Paths to the image files on the computer.
+     *
+     * @return If the import did not exit because of an error.
+     *
+     * @throws ServiceException   Cannot connect to OMERO.
+     * @throws AccessException    Cannot access data.
+     * @throws OMEROServerError   Server error.
+     * @throws IOException        Cannot read file.
+     * @throws ExecutionException A Facility can't be retrieved or instantiated.
+     */
+    public boolean importImages(GatewayWrapper client, int threads, String... paths)
+    throws ServiceException, OMEROServerError, AccessException, IOException, ExecutionException {
+        boolean success = importImages(client, data, threads, paths);
         refresh(client);
         return success;
     }

src/test/java/fr/igred/omero/SudoTest.java

@@ -109,6 +109,8 @@ void sudoImport() throws Exception {
 
         List<ImageWrapper> images = dataset.getImages(client3);
         assertEquals(1, images.size());
+        assertEquals(client3.getId(), images.get(0).getOwner().getId());
+        assertEquals(6L, images.get(0).getGroupId());
 
         client4.delete(images.get(0));
         client4.delete(dataset);

src/test/java/fr/igred/omero/repository/ImageImportTest.java

@@ -49,7 +49,8 @@ void testImportImage() throws Exception {
 
         DatasetWrapper dataset = client.getDataset(DATASET2.id);
 
-        boolean imported = dataset.importImages(client, f1.getAbsolutePath(), f2.getAbsolutePath());
+        boolean imported = dataset.importImages(client, 2, f1.getAbsolutePath(), f2.getAbsolutePath());
+        client.closeImport();
 
         removeFile(f1);
         removeFile(f2);