Merge pull request #2210 from adamretter/hotfix/context-update-listener-concurrency

Fix a concurrency issue in ContextUpdateListener

Author: duncdrum

src/org/exist/storage/NotificationService.java

@@ -21,80 +21,81 @@
  */
 package org.exist.storage;
 
+import net.jcip.annotations.ThreadSafe;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.exist.dom.persistent.DocumentImpl;
 import org.exist.numbering.NodeId;
 
 import java.util.IdentityHashMap;
+import java.util.Map;
+
 import org.exist.dom.persistent.IStoredNode;
 
 /**
  * Global notification service for document updates. Other classes
  * can subscribe to this service to be notified of document modifications,
  * removals or additions.
- * 
- * @author wolf
  *
+ * @author wolf
  */
-public class NotificationService extends IdentityHashMap<UpdateListener, Object> implements BrokerPoolService {
+@ThreadSafe
+public class NotificationService implements BrokerPoolService {
+
+    private static final long serialVersionUID = -3629584664969740903L;
+    private static final Logger LOG = LogManager.getLogger(NotificationService.class);
+
+    private final Map<UpdateListener, Object> listeners = new IdentityHashMap<>();
+
+    public NotificationService() {
+        super();
+    }
 
-	private static final long serialVersionUID = -3629584664969740903L;
+    /**
+     * Subscribe an {@link UpdateListener} to receive notifications.
+     *
+     * @param listener
+     */
+    public synchronized void subscribe(final UpdateListener listener) {
+        listeners.put(listener, new Object());
+    }
 
-	private final static Logger LOG = LogManager.getLogger(NotificationService.class);
-	
-	public NotificationService() {
-		super();
-	}
-	
-	/**
-	 * Subscribe an {@link UpdateListener} to receive notifications.
-	 * 
-	 * @param listener
-	 */
-	public synchronized void subscribe(UpdateListener listener) {
-		put(listener, new Object());
-	}
-	
-	/**
-	 * Unsubscribe an {@link UpdateListener}.
-	 * 
-	 * @param listener
-	 */
-	public synchronized void unsubscribe(UpdateListener listener) {
-		final Object i = remove(listener);
-		if (i == null)
-			{throw new RuntimeException(hashCode() + " listener not found: " + listener.hashCode());}
+    /**
+     * Unsubscribe an {@link UpdateListener}.
+     *
+     * @param listener
+     */
+    public synchronized void unsubscribe(final UpdateListener listener) {
+        final Object i = listeners.remove(listener);
+        if (i == null) {
+            throw new RuntimeException(hashCode() + " listener not found: " + listener.hashCode());
+        }
         listener.unsubscribe();
     }
 
-	/**
-	 * Notify all subscribers that a document has been updated/removed or
-	 * a new document has been added.
-	 * 
-	 * @param document
-	 * @param event
-	 */
-	public synchronized void notifyUpdate(DocumentImpl document, int event) {
-		for (final UpdateListener listener : keySet()) {
-	        listener.documentUpdated(document, event);
-		}
-	}
+    /**
+     * Notify all subscribers that a document has been updated/removed or
+     * a new document has been added.
+     *
+     * @param document
+     * @param event
+     */
+    public synchronized void notifyUpdate(final DocumentImpl document, final int event) {
+        listeners.keySet().forEach(listener -> listener.documentUpdated(document, event));
+    }
 
     /**
-	 * Notify all subscribers that a node has been moved. Nodes may be moved during a
+     * Notify all subscribers that a node has been moved. Nodes may be moved during a
      * defragmentation run.
-	 */
-	public synchronized void notifyMove(NodeId oldNodeId, IStoredNode newNode) {
-		for (final UpdateListener listener : keySet()) {
-	        listener.nodeMoved(oldNodeId, newNode);
-		}
-	}
+     */
+    public synchronized void notifyMove(final NodeId oldNodeId, final IStoredNode newNode) {
+        listeners.keySet().forEach(listener -> listener.nodeMoved(oldNodeId, newNode));
+    }
 
-    public void debug() {
-		LOG.debug("Registered UpdateListeners:");
-		for (final UpdateListener listener : keySet()) {
-	        listener.debug();
-		}
-	}
+    public synchronized void debug() {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Registered UpdateListeners:");
+        }
+        listeners.keySet().forEach(UpdateListener::debug);
+    }
 }

src/org/exist/xquery/XQueryContext.java

@@ -30,6 +30,8 @@
 import java.nio.charset.Charset;
 import java.nio.file.Path;
 import java.util.*;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -50,6 +52,7 @@
 import com.evolvedbinary.j8fu.tuple.Tuple2;
 import com.ibm.icu.text.Collator;
 import net.jcip.annotations.Immutable;
+import net.jcip.annotations.ThreadSafe;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.exist.Database;
@@ -3121,38 +3124,57 @@ public void setSource(final Source source) {
         this.source = source;
     }
 
+    /**
+     * NOTE: the {@link #unsubscribe()} method can be called
+     * from {@link org.exist.storage.NotificationService#unsubscribe(UpdateListener)}
+     * by another thread, so this class needs to be thread-safe.
+     */
+    @ThreadSafe
     private static class ContextUpdateListener implements UpdateListener {
-        private final List<UpdateListener> listeners = new ArrayList<>();
+        /*
+         * We use Concurrent safe data structures here, so that we don't have
+         * to block any calling threads.
+         *
+         * The AtomicReference enables us to quickly clear the listeners
+         * in #unsubscribe() and maintain happens-before integrity whilst
+         * unsubcribing them. The CopyOnWriteArrayList allows
+         * us to add listeners whilst iterating over a snapshot
+         * of existing iterators in other methods.
+         */
+        private final AtomicReference<List<UpdateListener>> listeners = new AtomicReference<>(new CopyOnWriteArrayList<>());
 
         private void addListener(final UpdateListener listener) {
-            listeners.add(listener);
+            listeners.get().add(listener);
         }
 
         @Override
         public void documentUpdated(final DocumentImpl document, final int event) {
-            listeners.forEach(listener -> listener.documentUpdated(document, event));
+            listeners.get().forEach(listener -> listener.documentUpdated(document, event));
         }
 
         @Override
         public void unsubscribe() {
-            listeners.forEach(UpdateListener::unsubscribe);
-            listeners.clear();
+            List<UpdateListener> prev = listeners.get();
+            while (!listeners.compareAndSet(prev, new CopyOnWriteArrayList<>())) {
+                prev = listeners.get();
+            }
+
+            prev.forEach(UpdateListener::unsubscribe);
         }
 
         @Override
         public void nodeMoved(final NodeId oldNodeId, final NodeHandle newNode) {
-            listeners.forEach(listener -> listener.nodeMoved(oldNodeId, newNode));
+            listeners.get().forEach(listener -> listener.nodeMoved(oldNodeId, newNode));
         }
 
         @Override
         public void debug() {
             if (LOG.isDebugEnabled()) {
-                LOG.debug(String.format("XQueryContext: %s document update listeners", listeners.size()));
+                LOG.debug(String.format("XQueryContext: %s document update listeners", listeners.get().size()));
             }
 
-            listeners.forEach(UpdateListener::debug);
+            listeners.get().forEach(UpdateListener::debug);
         }
-
     }
 
     private final List<CleanupTask> cleanupTasks = new ArrayList<>();