chore: consistent reset of pagedquery counters

chore: null guard ConcurrentHashSet clear
chore: use idisposable to reset exclusion list in tests

Author: rvazarkar

src/CommonLib/ConcurrentHashSet.cs

@@ -54,12 +54,11 @@ public IEnumerable<string> Values() {
     }
 
     public void Clear() {
-        _backingDictionary.Clear();
+        _backingDictionary?.Clear();
     }
 
     public void Dispose() {
         _backingDictionary = null;
         GC.SuppressFinalize(this);
     }
-    
-}
+}

src/CommonLib/LdapConnectionPool.cs

@@ -286,8 +286,6 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> PagedQuery(LdapQuery
                     if (response != null) {
                         pageResponse = (PageResultResponseControl)response.Controls
                             .Where(x => x is PageResultResponseControl).DefaultIfEmpty(null).FirstOrDefault();
-                        // Reset retry counter on a successful response so the next page gets a fresh budget
-                        queryRetryCount = 0;
                     } else if (queryRetryCount == MaxRetries) {
                         _metric.Observe(LdapMetricDefinitions.FailedRequests, 1,
                             new LabelValues([nameof(LdapConnectionPool), _poolIdentifier]));
@@ -413,8 +411,9 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> PagedQuery(LdapQuery
                     continue;
                 }
 
-                // Reset busy retry count after a successfully delivered page so each page starts with a fresh budget
+                // Reset busy and query retry count after a successfully delivered page so each page starts with a fresh budget
                 busyRetryCount = 0;
+                queryRetryCount = 0;
 
                 foreach (SearchResultEntry entry in response.Entries) {
                     if (cancellationToken.IsCancellationRequested) {
@@ -581,6 +580,8 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(string distinguish
                      * If we get a busy error, we want to do an exponential backoff, but maintain the current connection.
                      * The expectation is that given enough time, the server should stop being busy and service our query appropriately.
                      */
+                    _metric.Observe(LdapMetricDefinitions.FailedRequests, 1,
+                        new LabelValues([nameof(LdapConnectionPool), _poolIdentifier]));
                     busyRetryCount++;
                     _log.LogDebug("RangedRetrieval - Executing busy backoff for query {Info} (Attempt {Count})",
                         queryParameters.GetQueryInfo(), busyRetryCount);
@@ -591,6 +592,8 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(string distinguish
                     /*
                      * Treat a timeout as a busy error
                      */
+                    _metric.Observe(LdapMetricDefinitions.FailedRequests, 1,
+                        new LabelValues([nameof(LdapConnectionPool), _poolIdentifier]));
                     busyRetryCount++;
                     _log.LogDebug("RangedRetrieval - Timeout: Executing busy backoff for query {Info} (Attempt {Count})",
                         queryParameters.GetQueryInfo(), busyRetryCount);

test/unit/LdapConnectionPoolTest.cs

@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Concurrent;
 using System.DirectoryServices.Protocols;
 using System.Reflection;
@@ -7,8 +8,12 @@
 using SharpHoundCommonLib;
 using Xunit;
 
-public class LdapConnectionPoolTest
+public class LdapConnectionPoolTest : IDisposable
 {
+    public void Dispose() {
+        ResetExclusionDomain();
+    }
+
     private static void AddExclusionDomain(string identifier) {
         var excludedDomainsField = typeof(LdapConnectionPool)
             .GetField("ExcludedDomains", BindingFlags.Static | BindingFlags.NonPublic);
@@ -163,4 +168,4 @@ public void LdapConnectionPool_ReleaseConnection_NonGlobalCatalog_RoutesToConnec
 
         pool.Dispose();
     }
-}
+}