From c6aa4b08b096f8244cffc3d2f6c9bbe83402c02d 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.
---
 src/java/org/apache/cassandra/service/StorageService.java  | 5 +++++
 .../org/apache/cassandra/service/StorageServiceMBean.java  | 2 ++
 src/java/org/apache/cassandra/tools/NodeProbe.java         | 7 +++++--
 3 files changed, 12 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..00d33240c5 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;
@@ -482,11 +483,13 @@ 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();
+        Map<String, CompositeData> toppartitions = ssProxy.getToppartitions(samplers.stream().map(sampler -> sampler.name().toLowerCase()).collect(Collectors.toList()),
+            keyspaceFilters, tableFilters, duration, capacity, count);
 
+        Map<Sampler, CompositeData> result = Maps.newHashMap();
         for (Sampler sampler : samplers)
         {
-            result.put(sampler, ssProxy.getToppartitions(sampler.name(), keyspaceFilters, tableFilters, duration, capacity, count));
+            result.put(sampler, toppartitions.get(sampler.name().toLowerCase()));
         }
 
         return result;