Skip to content

Commit 7f4af7e

Browse files
wchoichoiwaiyiu
wchoi
authored andcommitted
Export clusterstats to jmx
1 parent 897118c commit 7f4af7e

File tree

5 files changed

+149
-4
lines changed

5 files changed

+149
-4
lines changed

gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import io.trino.gateway.ha.resource.LoginResource;
3333
import io.trino.gateway.ha.resource.PublicResource;
3434
import io.trino.gateway.ha.resource.TrinoResource;
35+
import io.trino.gateway.ha.router.BackendStateMBeanExporter;
3536
import io.trino.gateway.ha.security.AuthorizedExceptionMapper;
3637
import io.trino.gateway.proxyserver.ForProxy;
3738
import io.trino.gateway.proxyserver.ProxyRequestHandler;
@@ -125,6 +126,7 @@ public void configure(Binder binder)
125126
jaxrsBinder(binder).bind(AuthorizedExceptionMapper.class);
126127
binder.bind(ProxyHandlerStats.class).in(Scopes.SINGLETON);
127128
newExporter(binder).export(ProxyHandlerStats.class).withGeneratedName();
129+
binder.bind(BackendStateMBeanExporter.class).in(Scopes.SINGLETON);
128130
}
129131

130132
private static void addManagedApps(HaGatewayConfiguration configuration, Binder binder)

gateway-ha/src/main/java/io/trino/gateway/ha/clustermonitor/ClusterStatsObserver.java

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,18 +13,23 @@
1313
*/
1414
package io.trino.gateway.ha.clustermonitor;
1515

16+
import io.trino.gateway.ha.router.BackendStateMBeanExporter;
1617
import io.trino.gateway.ha.router.BackendStateManager;
1718

1819
import java.util.List;
1920

21+
import static java.util.Objects.requireNonNull;
22+
2023
public class ClusterStatsObserver
2124
implements TrinoClusterStatsObserver
2225
{
2326
private final BackendStateManager backendStateManager;
27+
private final BackendStateMBeanExporter backendStateMBeanExporter;
2428

25-
public ClusterStatsObserver(BackendStateManager backendStateManager)
29+
public ClusterStatsObserver(BackendStateManager backendStateManager, BackendStateMBeanExporter backendStateMBeanExporter)
2630
{
27-
this.backendStateManager = backendStateManager;
31+
this.backendStateManager = requireNonNull(backendStateManager, "backendStateManager is null");
32+
this.backendStateMBeanExporter = requireNonNull(backendStateMBeanExporter, "backendStateMBeanExporter is null");
2833
}
2934

3035
@Override
@@ -33,5 +38,6 @@ public void observe(List<ClusterStats> clustersStats)
3338
for (ClusterStats clusterStats : clustersStats) {
3439
backendStateManager.updateStates(clusterStats.clusterId(), clusterStats);
3540
}
41+
backendStateMBeanExporter.updateExport();
3642
}
3743
}

gateway-ha/src/main/java/io/trino/gateway/ha/module/ClusterStateListenerModule.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import io.trino.gateway.ha.clustermonitor.TrinoClusterStatsObserver;
2424
import io.trino.gateway.ha.config.HaGatewayConfiguration;
2525
import io.trino.gateway.ha.config.MonitorConfiguration;
26+
import io.trino.gateway.ha.router.BackendStateMBeanExporter;
2627
import io.trino.gateway.ha.router.BackendStateManager;
2728
import io.trino.gateway.ha.router.RoutingManager;
2829

@@ -46,11 +47,12 @@ public ClusterStateListenerModule(HaGatewayConfiguration config)
4647
@Singleton
4748
public List<TrinoClusterStatsObserver> getClusterStatsObservers(
4849
RoutingManager mgr,
49-
BackendStateManager backendStateManager)
50+
BackendStateManager backendStateManager,
51+
BackendStateMBeanExporter backendStateMBeanExporter)
5052
{
5153
return ImmutableList.<TrinoClusterStatsObserver>builder()
5254
.add(new HealthCheckObserver(mgr))
53-
.add(new ClusterStatsObserver(backendStateManager))
55+
.add(new ClusterStatsObserver(backendStateManager, backendStateMBeanExporter))
5456
.build();
5557
}
5658

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.gateway.ha.router;
15+
16+
import com.google.common.collect.ImmutableMap;
17+
import com.google.errorprone.annotations.concurrent.GuardedBy;
18+
import com.google.inject.Inject;
19+
import io.trino.gateway.ha.clustermonitor.ClusterStats;
20+
import jakarta.annotation.PostConstruct;
21+
import jakarta.annotation.PreDestroy;
22+
import org.weakref.jmx.MBeanExport;
23+
import org.weakref.jmx.MBeanExporter;
24+
import org.weakref.jmx.Managed;
25+
26+
import java.util.ArrayList;
27+
import java.util.HashMap;
28+
import java.util.List;
29+
import java.util.Map;
30+
31+
import static java.util.Objects.requireNonNull;
32+
33+
public class BackendStateMBeanExporter
34+
{
35+
@GuardedBy("this")
36+
private final List<MBeanExport> mbeanExports = new ArrayList<>();
37+
38+
private final MBeanExporter exporter;
39+
private final BackendStateManager backendStateManager;
40+
private final Map<String, ClusterStatsJMX> backendStates = new HashMap<>();
41+
42+
@Inject
43+
public BackendStateMBeanExporter(MBeanExporter exporter, BackendStateManager backendStateManager)
44+
{
45+
this.exporter = requireNonNull(exporter, "exporter is null");
46+
this.backendStateManager = requireNonNull(backendStateManager, "backendStateManager is null");
47+
}
48+
49+
@PostConstruct
50+
public synchronized void updateExport()
51+
{
52+
for (ClusterStats clusterStats : backendStateManager.getAllBackendStates().values()) {
53+
String clusterId = clusterStats.clusterId();
54+
55+
if (backendStates.containsKey(clusterId)) {
56+
ClusterStatsJMX clusterStatsJMX = backendStates.get(clusterId);
57+
clusterStatsJMX.setFrom(clusterStats);
58+
}
59+
else {
60+
ClusterStatsJMX clusterStatsJMX = new ClusterStatsJMX();
61+
clusterStatsJMX.setFrom(clusterStats);
62+
mbeanExports.add(exporter.exportWithGeneratedName(
63+
clusterStatsJMX,
64+
ClusterStatsJMX.class,
65+
ImmutableMap.<String, String>builder()
66+
.put("name", "ClusterStats")
67+
.put("backend", clusterId)
68+
.build()));
69+
backendStates.put(clusterId, clusterStatsJMX);
70+
}
71+
}
72+
}
73+
74+
@PreDestroy
75+
public synchronized void unexport()
76+
{
77+
for (MBeanExport mbeanExport : mbeanExports) {
78+
mbeanExport.unexport();
79+
}
80+
mbeanExports.clear();
81+
}
82+
83+
public static class ClusterStatsJMX
84+
{
85+
private int numWorkerNodes;
86+
private boolean healthy;
87+
private String proxyTo;
88+
private String externalUrl;
89+
private String routingGroup;
90+
91+
public void setFrom(ClusterStats clusterStats)
92+
{
93+
numWorkerNodes = clusterStats.numWorkerNodes();
94+
healthy = clusterStats.healthy();
95+
proxyTo = clusterStats.proxyTo();
96+
externalUrl = clusterStats.externalUrl();
97+
routingGroup = clusterStats.routingGroup();
98+
}
99+
100+
@Managed
101+
public int getNumWorkerNodes()
102+
{
103+
return numWorkerNodes;
104+
}
105+
106+
@Managed
107+
public boolean isHealthy()
108+
{
109+
return healthy;
110+
}
111+
112+
@Managed
113+
public String getProxyTo()
114+
{
115+
return proxyTo;
116+
}
117+
118+
@Managed
119+
public String getExternalUrl()
120+
{
121+
return externalUrl;
122+
}
123+
124+
@Managed
125+
public String getRoutingGroup()
126+
{
127+
return routingGroup;
128+
}
129+
}
130+
}

gateway-ha/src/main/java/io/trino/gateway/ha/router/BackendStateManager.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@ public ClusterStats getBackendState(ProxyBackendConfiguration backend)
3434
return clusterStats.getOrDefault(name, ClusterStats.builder(name).build());
3535
}
3636

37+
public Map<String, ClusterStats> getAllBackendStates()
38+
{
39+
return clusterStats;
40+
}
41+
3742
public void updateStates(String clusterId, ClusterStats stats)
3843
{
3944
clusterStats.put(clusterId, stats);

0 commit comments

Comments
 (0)