[#1175] Restart CTS persistent search after connection loss

openam-core/src/main/java/com/iplanet/services/ldap/event/LDAPv3PersistentSearch.java

@@ -12,6 +12,7 @@
  * information: "Portions copyright [year] [name of copyright owner]".
  *
  * Copyright 2015-2016 ForgeRock AS.
+ * Portions Copyright 2026 3A Systems, LLC.
  */
 
 package com.iplanet.services.ldap.event;
@@ -54,6 +55,7 @@
 import org.forgerock.opendj.ldap.responses.SearchResultEntry;
 import org.forgerock.opendj.ldap.responses.SearchResultReference;
 import org.forgerock.util.annotations.VisibleForTesting;
+import org.forgerock.util.promise.ExceptionHandler;
 
 import com.sun.identity.common.GeneralTaskRunnable;
 import com.sun.identity.common.SystemTimerPool;
@@ -73,6 +75,7 @@ public abstract class LDAPv3PersistentSearch<T, H> {
     private static final boolean CHANGES_ONLY = true;
     private static final boolean RETURN_CONTROLS = true;
     private static final boolean IS_CRITICAL = true;
+    private static final long ERROR_LOG_INTERVAL = 60000;
     private static final List<String> AD_DEFAULT_ATTRIBUTES = Collections.unmodifiableList(Arrays.asList(
             AD_IS_DELETED_ATTR, AD_WHEN_CHANGED_ATTR, AD_WHEN_CREATED_ATTR));
 
@@ -88,6 +91,7 @@ public abstract class LDAPv3PersistentSearch<T, H> {
     private LdapPromise<Result> futureResult;
     private PersistentSearchMode mode;
     private RetryTask retryTask;
+    private volatile long lastErrorLogged = 0;
 
     private enum PersistentSearchMode {
         STANDARD, AD, NONE
@@ -184,6 +188,20 @@ private void logError(Throwable t) {
         DEBUG.error("An error occurred while trying to initiate persistent search connection", t);
     }
 
+    /**
+     * Logs a failure of a running or restarting search at most once per {@link #ERROR_LOG_INTERVAL}. A directory that
+     * keeps refusing the search fails it again on every restart, so every retry interval would log it otherwise.
+     */
+    private void logRestartError(String message, Throwable t) {
+        long now = currentTimeMillis();
+        if (now - lastErrorLogged > ERROR_LOG_INTERVAL) {
+            lastErrorLogged = now;
+            DEBUG.error(message, t);
+        } else {
+            DEBUG.message(message, t);
+        }
+    }
+
     private void startSearch(Connection conn) throws LdapException {
         if (mode == null) {
             detectPersistentSearchMode(conn);
@@ -224,6 +242,7 @@ private void startSearch(Connection conn) throws LdapException {
         //cache, while PS was stopped, those gets cleared out and we start with a clean cache.
         clearCaches();
         futureResult = conn.searchAsync(searchRequest, null, new PersistentSearchResultHandler());
+        futureResult.thenOnException(new PersistentSearchExceptionHandler());
     }
 
     /**
@@ -365,10 +384,27 @@ public boolean handleReference(SearchResultReference reference) {
 
     }
 
+    /**
+     * Restarts the persistent search once the directory terminates it, e.g. because the connection was lost.
+     */
+    private class PersistentSearchExceptionHandler implements ExceptionHandler<LdapException> {
+
+        @Override
+        public void handleException(LdapException error) {
+            if (!shutdown) {
+                logRestartError("An error occurred while executing persistent search", error);
+                DEBUG.message("Restarting persistent search. Some changes may have been missed in the interim.");
+                clearCaches();
+                restartSearch();
+            } else {
+                DEBUG.message("Persistent search has been cancelled", error);
+            }
+        }
+    }
+
     private class RetryTask extends GeneralTaskRunnable {
 
         private long runPeriod;
-        private long lastLogged = 0;
 
         public RetryTask() {
             runPeriod = retryInterval;
@@ -391,18 +427,18 @@ public long getRunPeriod() {
         }
 
         public void run() {
+            //the failure handler may have scheduled this task while the search was being stopped
+            if (shutdown) {
+                runPeriod = -1;
+                return;
+            }
             try {
                 conn = factory.create();
                 startSearch(conn);
-                //everything seems to work, let's disable retryTask and reset the debug limit
+                //everything seems to work, let's disable retryTask
                 runPeriod = -1;
-                lastLogged = 0;
             } catch (DataLayerException | LdapException ex) {
-                long now = currentTimeMillis();
-                if (now - lastLogged > 60000) {
-                    DEBUG.error("Unable to start persistent search: " + ex.getMessage());
-                    lastLogged = now;
-                }
+                logRestartError("Unable to start persistent search: " + ex.getMessage(), null);
                 IOUtils.closeIfNotNull(conn);
             }
         }

openam-core/src/main/java/org/forgerock/openam/sm/datalayer/impl/ldap/CTSDJLDAPv3PersistentSearch.java

@@ -12,13 +12,15 @@
 * information: "Portions copyright [year] [name of copyright owner]".
 *
 * Copyright 2016 ForgeRock AS.
+* Portions Copyright 2026 3A Systems, LLC.
 */
 package org.forgerock.openam.sm.datalayer.impl.ldap;
 
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.forgerock.openam.cts.continuous.ChangeType;
 import org.forgerock.openam.cts.continuous.ContinuousQuery;
@@ -44,6 +46,7 @@ public class CTSDJLDAPv3PersistentSearch extends LDAPv3PersistentSearch<Continuo
         implements ContinuousQuery {
 
     private final SearchResultEntryHandler resultEntryHandler = new PSearchResultEntryHandler();
+    private final AtomicBoolean searchStarted = new AtomicBoolean(false);
 
     /**
      * Generate a new CTSDJLDAPv3PersistentSearch, providing the connection factory to use to produce
@@ -63,6 +66,10 @@ public CTSDJLDAPv3PersistentSearch(int retry, DN searchBaseDN, Filter searchFilt
 
     @Override
     protected void clearCaches() {
+        // The first call comes from the initial start of the search, when no change can have been missed yet.
+        if (searchStarted.compareAndSet(false, true)) {
+            return;
+        }
         for (ContinuousQueryListener listener : getListeners().keySet()) {
             listener.connectionLost();
         }

openam-core/src/test/java/org/forgerock/openam/sm/datalayer/impl/ldap/CTSDJLDAPv3PersistentSearchTest.java

@@ -12,19 +12,34 @@
 * information: "Portions copyright [year] [name of copyright owner]".
 *
 * Copyright 2016 ForgeRock AS.
+* Portions Copyright 2026 3A Systems, LLC.
 */
 package org.forgerock.openam.sm.datalayer.impl.ldap;
 
 import static org.assertj.core.api.Assertions.*;
+import static org.forgerock.opendj.ldap.LdapException.newLdapException;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.*;
 
 import java.util.HashSet;
 
 import org.forgerock.openam.cts.continuous.ContinuousQueryListener;
 import org.forgerock.openam.sm.datalayer.api.ConnectionFactory;
+import org.forgerock.opendj.ldap.Connection;
 import org.forgerock.opendj.ldap.DN;
 import org.forgerock.opendj.ldap.Filter;
+import org.forgerock.opendj.ldap.LdapException;
+import org.forgerock.opendj.ldap.ResultCode;
+import org.forgerock.opendj.ldap.SearchResultHandler;
 import org.forgerock.opendj.ldap.SearchScope;
+import org.forgerock.opendj.ldap.controls.PersistentSearchRequestControl;
+import org.forgerock.opendj.ldap.requests.SearchRequest;
+import org.forgerock.opendj.ldap.responses.Responses;
+import org.forgerock.opendj.ldap.responses.Result;
+import org.forgerock.opendj.ldap.spi.LdapPromises;
+import org.forgerock.util.promise.PromiseImpl;
+import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -42,22 +57,132 @@ public void theSetup() { //you need this
                 SearchScope.WHOLE_SUBTREE, mockConnectionFactory);
     }
 
+    @AfterMethod
+    public void tearDown() {
+        pSearch.stopQuery();
+    }
+
+    @Test
+    public void shouldNotInformListenersConnectionDownOnInitialStart() throws Exception {
+        //given
+        ContinuousQueryListener mockListener = mock(ContinuousQueryListener.class);
+        pSearch.addContinuousQueryListener(mockListener);
+        givenDirectoryAcceptsPersistentSearch(PromiseImpl.<Result, LdapException>create());
+
+        //when
+        pSearch.startQuery();
+
+        //then
+        verify(mockListener, never()).connectionLost();
+    }
+
     @Test
-    public void shouldInformListenersConnectionDown() {
+    public void shouldInformListenersConnectionDownWhenSearchFails() throws Exception {
         //given
         ContinuousQueryListener mockListener = mock(ContinuousQueryListener.class);
         ContinuousQueryListener mockListener2 = mock(ContinuousQueryListener.class);
         pSearch.addContinuousQueryListener(mockListener);
         pSearch.addContinuousQueryListener(mockListener2);
+        PromiseImpl<Result, LdapException> searchResult = PromiseImpl.create();
+        givenDirectoryAcceptsPersistentSearch(searchResult);
+        pSearch.startQuery();
 
         //when
-        pSearch.clearCaches();
+        searchResult.handleException(newLdapException(ResultCode.UNAVAILABLE));
 
         //then
         verify(mockListener, times(1)).connectionLost();
         verify(mockListener2, times(1)).connectionLost();
     }
 
+    @Test
+    public void shouldRestartSearchWhenSearchFails() throws Exception {
+        //given
+        pSearch = new CTSDJLDAPv3PersistentSearch(1000, DN.rootDN(), Filter.alwaysFalse(),
+                SearchScope.WHOLE_SUBTREE, mockConnectionFactory);
+        pSearch.addContinuousQueryListener(mock(ContinuousQueryListener.class));
+        PromiseImpl<Result, LdapException> searchResult = PromiseImpl.create();
+        givenDirectoryAcceptsPersistentSearch(searchResult);
+        pSearch.startQuery();
+
+        //when
+        searchResult.handleException(newLdapException(ResultCode.UNAVAILABLE));
+
+        //then
+        verify(mockConnectionFactory, timeout(5000).times(2)).create();
+    }
+
+    @Test
+    public void shouldRestartSearchWhenDirectoryRefusesSearch() throws Exception {
+        //given
+        pSearch = new CTSDJLDAPv3PersistentSearch(1000, DN.rootDN(), Filter.alwaysFalse(),
+                SearchScope.WHOLE_SUBTREE, mockConnectionFactory);
+        pSearch.addContinuousQueryListener(mock(ContinuousQueryListener.class));
+        PromiseImpl<Result, LdapException> searchResult = PromiseImpl.create();
+        givenDirectoryAcceptsPersistentSearch(searchResult);
+        pSearch.startQuery();
+
+        //when
+        searchResult.handleException(newLdapException(ResultCode.INSUFFICIENT_ACCESS_RIGHTS));
+
+        //then - the access rights may be fixed without restarting AM, so the search keeps being restarted
+        verify(mockConnectionFactory, timeout(5000).times(2)).create();
+    }
+
+    @Test
+    public void shouldInformListenersWhenSearchIsReestablished() throws Exception {
+        //given
+        pSearch = new CTSDJLDAPv3PersistentSearch(1000, DN.rootDN(), Filter.alwaysFalse(),
+                SearchScope.WHOLE_SUBTREE, mockConnectionFactory);
+        ContinuousQueryListener mockListener = mock(ContinuousQueryListener.class);
+        pSearch.addContinuousQueryListener(mockListener);
+        PromiseImpl<Result, LdapException> failedSearch = PromiseImpl.create();
+        givenDirectoryAcceptsPersistentSearch(failedSearch, PromiseImpl.<Result, LdapException>create());
+        pSearch.startQuery();
+
+        //when
+        failedSearch.handleException(newLdapException(ResultCode.UNAVAILABLE));
+
+        //then - one call when the search fails, one when it is re-established
+        verify(mockConnectionFactory, timeout(5000).times(2)).create();
+        verify(mockListener, timeout(5000).times(2)).connectionLost();
+    }
+
+    @Test
+    public void shouldNotRestartSearchWhenStopped() throws Exception {
+        //given
+        pSearch = new CTSDJLDAPv3PersistentSearch(1000, DN.rootDN(), Filter.alwaysFalse(),
+                SearchScope.WHOLE_SUBTREE, mockConnectionFactory);
+        ContinuousQueryListener mockListener = mock(ContinuousQueryListener.class);
+        pSearch.addContinuousQueryListener(mockListener);
+        PromiseImpl<Result, LdapException> searchResult = PromiseImpl.create();
+        givenDirectoryAcceptsPersistentSearch(searchResult);
+        pSearch.startQuery();
+        pSearch.stopQuery();
+
+        //when
+        searchResult.handleException(newLdapException(ResultCode.CLIENT_SIDE_USER_CANCELLED));
+
+        //then
+        verify(mockListener, never()).connectionLost();
+        verify(mockConnectionFactory, after(2000).times(1)).create();
+    }
+
+    private void givenDirectoryAcceptsPersistentSearch(PromiseImpl<Result, LdapException> searchResult)
+            throws Exception {
+        givenDirectoryAcceptsPersistentSearch(searchResult, searchResult);
+    }
+
+    private void givenDirectoryAcceptsPersistentSearch(PromiseImpl<Result, LdapException> firstSearchResult,
+            PromiseImpl<Result, LdapException> laterSearchResult) throws Exception {
+        Connection connection = mock(Connection.class);
+        given(mockConnectionFactory.create()).willReturn(connection);
+        given(connection.searchSingleEntry(any(SearchRequest.class))).willReturn(
+                Responses.newSearchResultEntry("").addAttribute("supportedControl", PersistentSearchRequestControl.OID));
+        given(connection.searchAsync(any(SearchRequest.class), any(), any(SearchResultHandler.class)))
+                .willReturn(LdapPromises.asPromise(firstSearchResult), LdapPromises.asPromise(laterSearchResult));
+    }
+
     @Test
     public void shouldStopAllListenersAndStopSearch() {
         //given