From a40e1f340e1233ab028b5e69b2504fb9d273c86c Mon Sep 17 00:00:00 2001
From: Piotr Wojtczak <piotr.wojtczak@scylladb.com>
Date: Mon, 10 May 2021 12:01:03 +0200
Subject: [PATCH] toppartitions: Fix toppartitions to only jmx once

In line with previous API, nodetool called jmx's toppartitions
for each sampler separately. A better way to do this is to call
it once and then pick results for the samplers we need.
---
 .../cassandra/service/StorageService.java     |  5 ++++
 .../service/StorageServiceMBean.java          |  2 ++
 .../org/apache/cassandra/tools/NodeProbe.java | 26 +++++++++++++++++--
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/java/org/apache/cassandra/service/StorageService.java b/src/java/org/apache/cassandra/service/StorageService.java
index cd7541ce0f..e45a9413e1 100644
--- a/src/java/org/apache/cassandra/service/StorageService.java
+++ b/src/java/org/apache/cassandra/service/StorageService.java
@@ -5205,4 +5205,9 @@ public CompositeData getToppartitions(String sampler, List<String> keyspaceFilte
     {
         return null;
     }
+
+    public Map<String, CompositeData> getToppartitions(List<String> samplers, List<String> keyspaceFilters, List<String> tableFilters, int duration, int capacity, int count) throws OpenDataException
+    {
+        return null;
+    }
 }
diff --git a/src/java/org/apache/cassandra/service/StorageServiceMBean.java b/src/java/org/apache/cassandra/service/StorageServiceMBean.java
index 25b5b875f6..2f40cdaa24 100644
--- a/src/java/org/apache/cassandra/service/StorageServiceMBean.java
+++ b/src/java/org/apache/cassandra/service/StorageServiceMBean.java
@@ -675,6 +675,8 @@ public interface StorageServiceMBean extends NotificationEmitter
 
     public CompositeData getToppartitions(String sampler, List<String> keyspaceFilters, List<String> tableFilters, int duration, int capacity, int count) throws OpenDataException;
 
+    public Map<String, CompositeData> getToppartitions(List<String> samplers, List<String> keyspaceFilters, List<String> tableFilters, int duration, int capacity, int count) throws OpenDataException;
+
     /**
      * Resume bootstrap streaming when there is failed data streaming.
      *
diff --git a/src/java/org/apache/cassandra/tools/NodeProbe.java b/src/java/org/apache/cassandra/tools/NodeProbe.java
index 3fe00a6199..adf0eaa5b7 100644
--- a/src/java/org/apache/cassandra/tools/NodeProbe.java
+++ b/src/java/org/apache/cassandra/tools/NodeProbe.java
@@ -44,6 +44,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
 
 import javax.management.InstanceNotFoundException;
 import javax.management.JMX;
@@ -483,10 +484,31 @@ public Map<Sampler, CompositeData> getPartitionSample(String ks, String cf, int
     public Map<Sampler, CompositeData> getToppartitions(List<String> keyspaceFilters, List<String> tableFilters, int capacity, int duration, int count, List<Sampler> samplers) throws OpenDataException
     {
         Map<Sampler, CompositeData> result = Maps.newHashMap();
+        
+        try
+        {
+            Class[] cArg = new Class[6];
+            cArg[0] = cArg[1] = cArg[2] = (Class<List<String>>) Collections.<String>emptyList().getClass();
+            cArg[3] = cArg[4] = cArg[5] = int.class;
+
+            // make sure that JMX has the newer API, throws NoSuchMethodException if not
+            ssProxy.getClass().getMethod("getToppartitions", cArg);
 
-        for (Sampler sampler : samplers)
+            Map<String, CompositeData> toppartitions = ssProxy.getToppartitions(samplers.stream().map(sampler -> sampler.name().toLowerCase()).collect(Collectors.toList()),
+                keyspaceFilters, tableFilters, duration, capacity, count);
+
+            for (Sampler sampler : samplers)
+            {
+                result.put(sampler, toppartitions.get(sampler.name().toLowerCase()));
+            }
+        }
+        catch (NoSuchMethodException e) 
         {
-            result.put(sampler, ssProxy.getToppartitions(sampler.name(), keyspaceFilters, tableFilters, duration, capacity, count));
+            // fall back to the old JMX API
+            for (Sampler sampler : samplers) 
+            {
+                result.put(sampler, ssProxy.getToppartitions(sampler.name(), keyspaceFilters, tableFilters, duration, capacity, count));
+            }
         }
 
         return result;