Expose /v1/integrations/trinoGateway/metrics for metrics used by trino gateway

Author: andythsu

core/trino-main/src/main/java/io/trino/memory/LocalMemoryManager.java

@@ -17,8 +17,10 @@
 import com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
 import com.google.inject.Inject;
+import com.sun.management.UnixOperatingSystemMXBean;
 import io.airlift.units.DataSize;
 
+import java.lang.management.ManagementFactory;
 import java.util.concurrent.TimeUnit;
 
 import static com.google.common.base.Verify.verify;
@@ -63,7 +65,7 @@ private void validateHeapHeadroom(NodeMemoryConfig config, long availableMemory)
 
     public MemoryInfo getInfo()
     {
-        return new MemoryInfo(AVAILABLE_PROCESSORS.get(), memoryPool.getInfo());
+        return new MemoryInfo(AVAILABLE_PROCESSORS.get(), ((UnixOperatingSystemMXBean) ManagementFactory.getOperatingSystemMXBean()).getCpuLoad(), memoryPool.getInfo());
     }
 
     public MemoryPool getMemoryPool()

core/trino-main/src/main/java/io/trino/memory/MemoryInfo.java

@@ -23,14 +23,17 @@
 public class MemoryInfo
 {
     private final int availableProcessors;
+    private final double systemCpuLoad;
     private final MemoryPoolInfo pool;
 
     @JsonCreator
     public MemoryInfo(
             @JsonProperty("availableProcessors") int availableProcessors,
+            @JsonProperty("systemCpuLoad") double systemCpuLoad,
             @JsonProperty("pool") MemoryPoolInfo pool)
     {
         this.availableProcessors = availableProcessors;
+        this.systemCpuLoad = systemCpuLoad;
         this.pool = requireNonNull(pool, "pool is null");
     }
 
@@ -40,6 +43,12 @@ public int getAvailableProcessors()
         return availableProcessors;
     }
 
+    @JsonProperty
+    public double getSystemCpuLoad()
+    {
+        return systemCpuLoad;
+    }
+
     @JsonProperty
     public MemoryPoolInfo getPool()
     {
@@ -51,6 +60,7 @@ public String toString()
     {
         return toStringHelper(this)
                 .add("availableProcessors", availableProcessors)
+                .add("systemCpuLoad", systemCpuLoad)
                 .add("pool", pool)
                 .toString();
     }

core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java

@@ -242,6 +242,8 @@ protected void setup(Binder binder)
 
         newExporter(binder).export(ClusterMemoryManager.class).withGeneratedName();
 
+        jaxrsBinder(binder).bind(GatewayResource.class);
+
         // node partitioning manager
         binder.bind(NodePartitioningManager.class).in(Scopes.SINGLETON);
 

core/trino-main/src/main/java/io/trino/server/GatewayResource.java

@@ -0,0 +1,70 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.server;
+
+import com.google.inject.Inject;
+import io.trino.memory.ClusterMemoryManager;
+import io.trino.memory.MemoryInfo;
+import io.trino.server.security.ResourceSecurity;
+import io.trino.spi.memory.MemoryPoolInfo;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+
+import java.util.Map;
+import java.util.Optional;
+
+import static io.trino.server.security.ResourceSecurity.AccessType.MANAGEMENT_READ;
+import static java.util.Objects.requireNonNull;
+
+@Path("/v1/integrations/gateway")
+public class GatewayResource
+{
+    private final ClusterMemoryManager clusterMemoryManager;
+
+    @Inject
+    public GatewayResource(ClusterMemoryManager clusterMemoryManager)
+    {
+        this.clusterMemoryManager = requireNonNull(clusterMemoryManager, "clusterMemoryManager is null");
+    }
+
+    @GET
+    @Path("clusterMetrics")
+    @ResourceSecurity(MANAGEMENT_READ)
+    public ClusterMetrics getClusterMetrics()
+    {
+        Map<String, Optional<MemoryInfo>> memoryInfo = clusterMemoryManager.getAllNodesMemoryInfo();
+        long totalFreeBytes = memoryInfo
+                .values()
+                .stream()
+                .flatMap(Optional::stream)
+                .map(MemoryInfo::getPool)
+                .mapToLong(MemoryPoolInfo::getFreeBytes)
+                .sum();
+        double aggregatedSystemLoad = memoryInfo
+                .values()
+                .stream()
+                .flatMap(Optional::stream)
+                .mapToDouble(MemoryInfo::getSystemCpuLoad)
+                .sum();
+        return new ClusterMetrics(totalFreeBytes, aggregatedSystemLoad);
+    }
+
+    /**
+     * Represents the aggregated metrics from each node in Trino cluster
+     *
+     * @param totalFreeBytes the sum of free memory from each node in the cluster
+     * @param aggregatedSystemLoad the sum of system load from each node in the cluster
+     */
+    public record ClusterMetrics(long totalFreeBytes, double aggregatedSystemLoad) {}
+}

core/trino-main/src/test/java/io/trino/execution/scheduler/faulttolerant/BenchmarkBinPackingNodeAllocator.java

@@ -179,6 +179,7 @@ private MemoryInfo buildWorkerMemoryInfo(DataSize usedMemory, Map<TaskId, DataSi
         {
             return new MemoryInfo(
                     4,
+                    0,
                     new MemoryPoolInfo(
                             DataSize.of(64, GIGABYTE).toBytes(),
                             usedMemory.toBytes(),

core/trino-main/src/test/java/io/trino/execution/scheduler/faulttolerant/TestBinPackingNodeAllocator.java

@@ -131,6 +131,7 @@ private MemoryInfo buildWorkerMemoryInfo(DataSize usedMemory, Map<TaskId, DataSi
     {
         return new MemoryInfo(
                 4,
+                0,
                 new MemoryPoolInfo(
                         DataSize.of(64, GIGABYTE).toBytes(),
                         usedMemory.toBytes(),

core/trino-main/src/test/java/io/trino/execution/scheduler/faulttolerant/TestExponentialGrowthPartitionMemoryEstimator.java

@@ -279,6 +279,7 @@ private MemoryInfo buildWorkerMemoryInfo(DataSize usedMemory)
     {
         return new MemoryInfo(
                 4,
+                0,
                 new MemoryPoolInfo(
                         DataSize.of(64, GIGABYTE).toBytes(),
                         usedMemory.toBytes(),

core/trino-main/src/test/java/io/trino/memory/LowMemoryKillerTestingUtils.java

@@ -73,7 +73,7 @@ static List<MemoryInfo> toNodeMemoryInfoList(long memoryPoolMaxBytes, Map<String
                     ImmutableMap.of(),
                     tasksMemoryInfoForNode(entry.getKey().getNodeIdentifier(), tasks),
                     ImmutableMap.of());
-            result.add(new MemoryInfo(7, memoryPoolInfo));
+            result.add(new MemoryInfo(7, 0, memoryPoolInfo));
         }
         return result.build();
     }