Service, tuning review (#1279)

scottf · web-flow · commit f60390e8a190 · 2025-02-20T12:17:07.000-05:00
* Service, tuning review

* Service, tuning review

* Service, tuning review
diff --git a/src/main/java/io/nats/service/EndpointContext.java b/src/main/java/io/nats/service/EndpointContext.java
@@ -3,10 +3,11 @@
 import io.nats.client.Connection;
 import io.nats.client.Dispatcher;
 import io.nats.client.Message;
-import io.nats.client.Subscription;
 import io.nats.client.support.DateTimeUtils;
 
+import java.time.Duration;
 import java.time.ZonedDateTime;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
@@ -19,12 +20,11 @@ class EndpointContext {
     private final ServiceMessageHandler handler;
     private final boolean recordStats;
     private final String qGroup;
+    private boolean running;
 
     private final boolean internalDispatcher;
     private final Dispatcher dispatcher;
 
-    private Subscription sub;
-
     private ZonedDateTime started;
     private String lastError;
     private final AtomicLong numRequests;
@@ -37,6 +37,7 @@ class EndpointContext {
         handler = se.getHandler();
         this.recordStats = !internalEndpoint;
         qGroup = internalEndpoint ? null : se.getQueueGroup();
+        running = false;
 
         if (se.getDispatcher() == null) {
             dispatcher = internalDispatcher;
@@ -53,15 +54,19 @@ class EndpointContext {
         started = DateTimeUtils.gmtNow();
     }
 
+    // this method does not need a lock because it is only
+    // called from Service start which is already locked
     void start() {
-        if (sub == null) {
+        if (!running) {
+            // we do not need to track the sub, since drain occurs through the dispatcher
             if (qGroup == null) {
                 dispatcher.subscribe(se.getSubject(), this::onMessage);
             }
             else {
                 dispatcher.subscribe(se.getSubject(), qGroup, this::onMessage);
             }
             started = DateTimeUtils.gmtNow();
+            running = true;
         }
     }
 
@@ -115,7 +120,7 @@ boolean isNotInternalDispatcher() {
         return !internalDispatcher;
     }
 
-    Subscription getSub() {
-        return sub;
+    CompletableFuture<Boolean> drain(Duration timeout) throws InterruptedException {
+        return dispatcher.drain(timeout);
     }
 }
diff --git a/src/main/java/io/nats/service/InfoResponse.java b/src/main/java/io/nats/service/InfoResponse.java
@@ -16,10 +16,7 @@
 import io.nats.client.support.JsonUtils;
 import io.nats.client.support.JsonValue;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
+import java.util.*;
 
 import static io.nats.client.support.ApiConstants.DESCRIPTION;
 import static io.nats.client.support.ApiConstants.ENDPOINTS;
@@ -42,15 +39,25 @@ public class InfoResponse extends ServiceResponse {
         this.endpoints = new ArrayList<>();
     }
 
-    InfoResponse addServiceEndpoint(ServiceEndpoint se) {
+    void addServiceEndpoint(ServiceEndpoint se) {
+        _addServiceEndpoint(se);
+        serialized.set(null);
+    }
+
+    void addServiceEndpoints(Collection<ServiceEndpoint> serviceEndpoints) {
+        for (ServiceEndpoint se : serviceEndpoints) {
+            _addServiceEndpoint(se);
+        }
+        serialized.set(null);
+    }
+
+    private void _addServiceEndpoint(ServiceEndpoint se) {
         endpoints.add(new Endpoint(
             se.getName(),
             se.getSubject(),
             se.getQueueGroup(),
             se.getMetadata()
         ));
-        serialized.set(null);
-        return this;
     }
 
     InfoResponse(byte[] jsonBytes) {
diff --git a/src/main/java/io/nats/service/Service.java b/src/main/java/io/nats/service/Service.java
@@ -21,10 +21,11 @@
 import java.time.Duration;
 import java.time.ZonedDateTime;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -46,7 +47,7 @@ public class Service {
 
     private final Connection conn;
     private final Duration drainTimeout;
-    private final ConcurrentMap<String, EndpointContext> serviceContexts;
+    private final ConcurrentHashMap<String, EndpointContext> serviceContexts;
     private final List<EndpointContext> discoveryContexts;
     private final List<Dispatcher> dInternals;
     private final PingResponse pingResponse;
@@ -70,9 +71,7 @@ public class Service {
         // set up the service contexts
         // ? do we need an internal dispatcher for any user endpoints !! addServiceEndpoint deals with it
         serviceContexts = new ConcurrentHashMap<>();
-        for (ServiceEndpoint se : b.serviceEndpoints.values()) {
-            addServiceEndpoint(se);
-        }
+        addServiceEndpoints(b.serviceEndpoints.values());
 
         Dispatcher dTemp = null;
         if (b.pingDispatcher == null || b.infoDispatcher == null || b.statsDispatcher == null) {
@@ -87,31 +86,41 @@ public class Service {
     }
 
     /**
-     * Adds a service endpoint to the list of service contexts and starts it if the service is running.
-     * @param se the service endpoint to be added
+     * Adds one or more service endpoint to the list of service contexts and starts it if the service is running.
+     * @param serviceEndpoints one or more service endpoints to be added
      */
-    public void addServiceEndpoint(ServiceEndpoint se) {
-        // do this first so it's available on start
-        infoResponse.addServiceEndpoint(se);
-
-        EndpointContext ctx;
-        if (se.getDispatcher() == null) {
-            Dispatcher dTemp = dInternals.isEmpty() ? null : dInternals.get(0);
-            if (dTemp == null) {
-                dTemp = conn.createDispatcher();
-                dInternals.add(dTemp);
-            }
-            ctx = new EndpointContext(conn, dTemp, false, se);
-        } else {
-            ctx = new EndpointContext(conn, null, false, se);
-        }
-        serviceContexts.put(se.getName(), ctx);
+    public void addServiceEndpoints(ServiceEndpoint... serviceEndpoints) {
+        addServiceEndpoints(Arrays.asList(serviceEndpoints));
+    }
 
-        // if the service is already started, start the newly added context
+    /**
+     * Adds all service endpoints to the list of service contexts and starts it if the service is running.
+     * @param serviceEndpoints service endpoints to be added
+     */
+    public void addServiceEndpoints(Collection<ServiceEndpoint> serviceEndpoints) {
         startStopLock.lock();
         try {
-            if (runningIndicator != null) {
-                ctx.start();
+            // do this first so it's available on start
+            infoResponse.addServiceEndpoints(serviceEndpoints);
+            for (ServiceEndpoint se : serviceEndpoints) {
+                EndpointContext ctx;
+                if (se.getDispatcher() == null) {
+                    Dispatcher dTemp = dInternals.isEmpty() ? null : dInternals.get(0);
+                    if (dTemp == null) {
+                        dTemp = conn.createDispatcher();
+                        dInternals.add(dTemp);
+                    }
+                    ctx = new EndpointContext(conn, dTemp, false, se);
+                }
+                else {
+                    ctx = new EndpointContext(conn, null, false, se);
+                }
+                serviceContexts.put(se.getName(), ctx);
+
+                // if the service is already started, start the newly added context
+                if (runningIndicator != null) {
+                    ctx.start();
+                }
             }
         }
         finally {
@@ -242,7 +251,7 @@ public void stop(boolean drain, Throwable t) {
                     for (EndpointContext c : serviceContexts.values()) {
                         if (c.isNotInternalDispatcher()) {
                             try {
-                                futures.add(c.getSub().drain(drainTimeout));
+                                futures.add(c.drain(drainTimeout));
                             }
                             catch (Exception e) { /* nothing I can really do, we are stopping anyway */ }
                         }
@@ -251,7 +260,7 @@ public void stop(boolean drain, Throwable t) {
                     for (EndpointContext c : discoveryContexts) {
                         if (c.isNotInternalDispatcher()) {
                             try {
-                                futures.add(c.getSub().drain(drainTimeout));
+                                futures.add(c.drain(drainTimeout));
                             }
                             catch (Exception e) { /* nothing I can really do, we are stopping anyway */ }
                         }
diff --git a/src/test/java/io/nats/service/ServiceTests.java b/src/test/java/io/nats/service/ServiceTests.java
@@ -260,7 +260,7 @@ public void testServiceWorkflow() throws Exception {
                     .handler(new ReverseHandler(serviceNc1))
                     .build();
 
-                service1.addServiceEndpoint(seRev1);
+                service1.addServiceEndpoints(seRev1);
 
                 for (int x = 0; x < requestCount; x++) {
                     verifyServiceExecution(clientNc, REVERSE_ENDPOINT_NAME, REVERSE_ENDPOINT_SUBJECT, null);
@@ -709,16 +709,7 @@ public void testDispatchers() throws Exception {
                 done.get(100, TimeUnit.MILLISECONDS);
 
                 dispatchers = getDispatchers(nc);
-                assertEquals(2, dispatchers.size()); // internal discovery was stopped
-                assertTrue(dispatchers.containsValue(dStats));
-                assertTrue(dispatchers.containsValue(dEnd));
-
-                nc.closeDispatcher(dStats);
-                nc.closeDispatcher(dEnd);
-                sleep(100); // no rush
-
-                dispatchers = getDispatchers(nc);
-                assertEquals(0, dispatchers.size());
+                assertEquals(0, dispatchers.size()); // stop() calls drain which closes dispatchers
 
                 se1 = ServiceEndpoint.builder()
                     .endpointName("dispatch")
@@ -839,7 +830,7 @@ public void testAddingEndpointAfterServiceBuilderConstruction() {
                 .drainTimeout(Duration.ofSeconds(1))
                 .build();
 
-        service.addServiceEndpoint(se);
+        service.addServiceEndpoints(se);
         assertEquals("desc", service.getDescription());
         assertEquals(Duration.ofSeconds(1), service.getDrainTimeout());
 
@@ -1375,7 +1366,8 @@ public void testServiceResponsesConstruction() {
         endMeta.put("foo", "bar");
         Endpoint ep = new Endpoint("endfoo", endMeta);
         ServiceEndpoint se = new ServiceEndpoint(ep, m -> {}, null);
-        InfoResponse ir1 = new InfoResponse("id", "name", "0.0.0", metadata, "desc").addServiceEndpoint(se);
+        InfoResponse ir1 = new InfoResponse("id", "name", "0.0.0", metadata, "desc");
+        ir1.addServiceEndpoint(se);
         InfoResponse ir2 = new InfoResponse(ir1.toJson().getBytes());
         validateApiInOutInfoResponse(ir1);
         validateApiInOutInfoResponse(ir2);
