Remove BackendState record

gympass · May 10, 2024 · 9db69d1 · 9db69d1
1 parent 53ec04b
commit 9db69d1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 23 deletions.
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/resource/GatewayWebAppResource.java b/gateway-ha/src/main/java/io/trino/gateway/ha/resource/GatewayWebAppResource.java
@@ -15,6 +15,7 @@
 
 import com.google.common.base.Strings;
 import com.google.inject.Inject;
+import io.trino.gateway.ha.clustermonitor.ClusterStats;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
 import io.trino.gateway.ha.domain.Result;
 import io.trino.gateway.ha.domain.TableData;
@@ -86,11 +87,10 @@ public Response getAllBackends()
     {
         List<ProxyBackendConfiguration> allBackends = gatewayBackendManager.getAllBackends();
         List<BackendResponse> data = allBackends.stream().map(b -> {
-            BackendStateManager.BackendState backendState = backendStateManager.getBackendState(b);
-            Map<String, Integer> state = backendState.state();
+            ClusterStats backendState = backendStateManager.getBackendState(b);
             BackendResponse backendResponse = new BackendResponse();
-            backendResponse.setQueued(state.get("QUEUED"));
-            backendResponse.setRunning(state.get("RUNNING"));
+            backendResponse.setQueued(backendState.queuedQueryCount());
+            backendResponse.setRunning(backendState.runningQueryCount());
             backendResponse.setName(b.getName());
             backendResponse.setProxyTo(b.getProxyTo());
             backendResponse.setActive(b.isActive());

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/resource/PublicResource.java b/gateway-ha/src/main/java/io/trino/gateway/ha/resource/PublicResource.java
@@ -59,7 +59,7 @@ public Response getBackend(@PathParam("name") String name)
     public Response getBackendState(@PathParam("name") String name)
     {
         return gatewayBackendManager.getBackendByName(name).map(backendStateManager::getBackendState)
-                .map(state -> Response.ok(state.state()).build())
+                .map(state -> Response.ok(state).build())
                 .orElseGet(() -> Response.status(404).build());
     }
 }
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/BackendStateManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/BackendStateManager.java
@@ -13,15 +13,12 @@
  */
 package io.trino.gateway.ha.router;
 
-import com.google.common.collect.ImmutableMap;
 import io.trino.gateway.ha.clustermonitor.ClusterStats;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
 
 import java.util.HashMap;
 import java.util.Map;
 
-import static java.util.Objects.requireNonNull;
-
 public class BackendStateManager
 {
     private final Map<String, ClusterStats> clusterStats;
@@ -31,27 +28,14 @@ public BackendStateManager()
         this.clusterStats = new HashMap<>();
     }
 
-    public BackendState getBackendState(ProxyBackendConfiguration backend)
+    public ClusterStats getBackendState(ProxyBackendConfiguration backend)
     {
         String name = backend.getName();
-        ClusterStats stats = clusterStats.getOrDefault(backend.getName(), ClusterStats.builder(name).build());
-        Map<String, Integer> state = new HashMap<>();
-        state.put("QUEUED", stats.queuedQueryCount());
-        state.put("RUNNING", stats.runningQueryCount());
-        return new BackendState(name, state);
+        return clusterStats.getOrDefault(name, ClusterStats.builder(name).build());
     }
 
     public void updateStates(String clusterId, ClusterStats stats)
     {
         clusterStats.put(clusterId, stats);
     }
-
-    public record BackendState(String name, Map<String, Integer> state)
-    {
-        public BackendState
-        {
-            requireNonNull(name, "name is null");
-            state = ImmutableMap.copyOf(requireNonNull(state, "state is null"));
-        }
-    }
 }