From 83fcab1d0f81081b5a307042ec1bdcbfa3f7d9cf Mon Sep 17 00:00:00 2001
From: Samarth Jain <samarth@apache.org>
Date: Wed, 17 Mar 2021 22:24:02 -0700
Subject: [PATCH] Improve performance of queries against SYSTEM.SEGMENT table.
 (#11008)

Size HashMap and HashSet appropriately. Perf analysis of the queries
revealed that over 25% of the query time was spent in resizing HashMap and HashSet
collections. Also, prevent the need to examine and authorize all resources when
AllowAllAuthorizer is the configured authorizer.
---
 .../org/apache/druid/server/security/AuthorizationUtils.java | 4 ++++
 .../org/apache/druid/sql/calcite/schema/DruidSchema.java     | 5 +++--
 .../org/apache/druid/sql/calcite/schema/SystemSchema.java    | 4 ++--
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java b/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java
index 64c6330f745..7017ac72a33 100644
--- a/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java
+++ b/server/src/main/java/org/apache/druid/server/security/AuthorizationUtils.java
@@ -263,6 +263,10 @@ public class AuthorizationUtils
       throw new ISE("No authorizer found with name: [%s].", authenticationResult.getAuthorizerName());
     }
 
+    if (authorizer instanceof AllowAllAuthorizer) {
+      return resources;
+    }
+
     final Map<ResourceAction, Access> resultCache = new HashMap<>();
     final Iterable<ResType> filteredResources = Iterables.filter(
         resources,
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/DruidSchema.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/DruidSchema.java
index 20138a401c7..30a1576afc6 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/DruidSchema.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/DruidSchema.java
@@ -702,13 +702,14 @@ public class DruidSchema extends AbstractSchema
 
   Map<SegmentId, AvailableSegmentMetadata> getSegmentMetadataSnapshot()
   {
-    final Map<SegmentId, AvailableSegmentMetadata> segmentMetadata = new HashMap<>();
     synchronized (lock) {
+      final Map<SegmentId, AvailableSegmentMetadata> segmentMetadata = Maps.newHashMapWithExpectedSize(
+          segmentMetadataInfo.values().stream().mapToInt(v -> v.size()).sum());
       for (TreeMap<SegmentId, AvailableSegmentMetadata> val : segmentMetadataInfo.values()) {
         segmentMetadata.putAll(val);
       }
+      return segmentMetadata;
     }
-    return segmentMetadata;
   }
 
   int getTotalSegments()
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java
index 40352eba107..0562e333314 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java
@@ -29,6 +29,7 @@ import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
 import com.google.common.net.HostAndPort;
 import com.google.common.util.concurrent.Futures;
 import com.google.inject.Inject;
@@ -90,7 +91,6 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -287,7 +287,7 @@ public class SystemSchema extends AbstractSchema
       // Coordinator.
       final Iterator<SegmentWithOvershadowedStatus> metadataStoreSegments = metadataView.getPublishedSegments();
 
-      final Set<SegmentId> segmentsAlreadySeen = new HashSet<>();
+      final Set<SegmentId> segmentsAlreadySeen = Sets.newHashSetWithExpectedSize(druidSchema.getTotalSegments());
 
       final FluentIterable<Object[]> publishedSegments = FluentIterable
           .from(() -> getAuthorizedPublishedSegments(metadataStoreSegments, root))