MINOR: Cache topic resolution in TopicIds set

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilder.java

@@ -327,13 +327,14 @@ public TargetAssignmentBuilder<T> removeMember(
      */
     public TargetAssignmentResult build() throws PartitionAssignorException {
         Map<String, MemberSubscriptionAndAssignmentImpl> memberSpecs = new HashMap<>();
+        TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(topicsImage);
 
         // Prepare the member spec for all members.
         members.forEach((memberId, member) ->
             memberSpecs.put(memberId, createMemberSubscriptionAndAssignment(
                 member,
                 targetAssignment.getOrDefault(memberId, Assignment.EMPTY),
-                topicsImage
+                topicResolver
             ))
         );
 
@@ -355,7 +356,7 @@ public TargetAssignmentResult build() throws PartitionAssignorException {
                 memberSpecs.put(memberId, createMemberSubscriptionAndAssignment(
                     updatedMemberOrNull,
                     assignment,
-                    topicsImage
+                    topicResolver
                 ));
             }
         });
@@ -420,12 +421,12 @@ private Assignment newMemberAssignment(
     static <T extends ModernGroupMember> MemberSubscriptionAndAssignmentImpl createMemberSubscriptionAndAssignment(
         T member,
         Assignment memberAssignment,
-        TopicsImage topicsImage
+        TopicIds.TopicResolver topicResolver
     ) {
         return new MemberSubscriptionAndAssignmentImpl(
             Optional.ofNullable(member.rackId()),
             Optional.ofNullable(member.instanceId()),
-            new TopicIds(member.subscribedTopicNames(), topicsImage),
+            new TopicIds(member.subscribedTopicNames(), topicResolver),
             memberAssignment
         );
     }

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/TopicIds.java

@@ -21,25 +21,169 @@
 import org.apache.kafka.image.TopicsImage;
 
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Set;
 
 /**
  * TopicIds is initialized with topic names (String) but exposes a Set of topic ids (Uuid) to the
- * user and performs the conversion lazily with TopicsImage.
+ * user and performs the conversion lazily with a TopicResolver backed by a TopicsImage.
  */
 public class TopicIds implements Set<Uuid> {
+    /**
+     * Converts between topic ids (Uuids) and topic names (Strings).
+     */
+    public interface TopicResolver {
+        /**
+         * @return The TopicsImage used by the resolver.
+         */
+        TopicsImage image();
+
+        /**
+         * Converts a topic id to a topic name.
+         *
+         * @param id The topic id.
+         * @return The topic name for the given topic id, or null if the topic does not exist.
+         */
+        String name(Uuid id);
+
+        /**
+         * Converts a topic name to a topic id.
+         *
+         * @param name The topic name.
+         * @return The topic id for the given topic name, or null if the topic does not exist.
+         */
+        Uuid id(String name);
+
+        /**
+         * Clears any cached data.
+         *
+         * Used for benchmarking purposes.
+         */
+        void clear();
+    }
+
+    /**
+     * A base implementation of TopicResolver.
+     *
+     * Provides an implementation of equals and hashCode based on the underlying TopicsImage.
+     */
+    public abstract static class BaseTopicResolver implements TopicResolver {
+        protected final TopicsImage image;
+
+        public BaseTopicResolver(
+            TopicsImage image
+        ) {
+            this.image = Objects.requireNonNull(image);
+        }
+
+        @Override
+        public final TopicsImage image() {
+            return image;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || !(o instanceof TopicResolver)) return false;
+
+            TopicResolver that = (TopicResolver) o;
+
+            return Objects.equals(image, that.image());
+        }
+
+        @Override
+        public int hashCode() {
+            return image.hashCode();
+        }
+    }
+
+    /**
+     * A TopicResolver without any caching.
+     */
+    public static class DefaultTopicResolver extends BaseTopicResolver {
+        public DefaultTopicResolver(
+            TopicsImage image
+        ) {
+            super(image);
+        }
+
+        @Override
+        public String name(Uuid id) {
+            TopicImage topic = image.getTopic(id);
+            if (topic == null) return null;
+            return topic.name();
+        }
+
+        @Override
+        public Uuid id(String name) {
+            TopicImage topic = image.getTopic(name);
+            if (topic == null) return null;
+            return topic.id();
+        }
+
+        @Override
+        public void clear() {}
+    }
+
+    /**
+     * A TopicResolver that caches results.
+     */
+    public static class CachedTopicResolver extends BaseTopicResolver {
+        private final Map<String, Uuid> topicIds = new HashMap<>();
+        private final Map<Uuid, String> topicNames = new HashMap<>();
+
+        public CachedTopicResolver(
+            TopicsImage image
+        ) {
+            super(image);
+        }
+
+        @Override
+        public String name(Uuid id) {
+            return topicNames.computeIfAbsent(id, __ -> {
+                TopicImage topic = image.getTopic(id);
+                if (topic == null) return null;
+                return topic.name();
+            });
+        }
+
+        @Override
+        public Uuid id(String name) {
+            return topicIds.computeIfAbsent(name, __ -> {
+                TopicImage topic = image.getTopic(name);
+                if (topic == null) return null;
+                return topic.id();
+            });
+        }
+
+        @Override
+        public void clear() {
+            this.topicNames.clear();
+            this.topicIds.clear();
+        }
+    }
+
     private final Set<String> topicNames;
-    private final TopicsImage image;
+    private final TopicResolver resolver;
 
     public TopicIds(
         Set<String> topicNames,
         TopicsImage image
     ) {
         this.topicNames = Objects.requireNonNull(topicNames);
-        this.image = Objects.requireNonNull(image);
+        this.resolver = new DefaultTopicResolver(image);
+    }
+
+    public TopicIds(
+        Set<String> topicNames,
+        TopicResolver resolver
+    ) {
+        this.topicNames = Objects.requireNonNull(topicNames);
+        this.resolver = Objects.requireNonNull(resolver);
     }
 
     @Override
@@ -56,24 +200,24 @@ public boolean isEmpty() {
     public boolean contains(Object o) {
         if (o instanceof Uuid) {
             Uuid topicId = (Uuid) o;
-            TopicImage topicImage = image.getTopic(topicId);
-            if (topicImage == null) return false;
-            return topicNames.contains(topicImage.name());
+            String topicName = resolver.name(topicId);
+            if (topicName == null) return false;
+            return topicNames.contains(topicName);
         }
         return false;
     }
 
     private static class TopicIdIterator implements Iterator<Uuid> {
         final Iterator<String> iterator;
-        final TopicsImage image;
+        final TopicResolver resolver;
         private Uuid next = null;
 
         private TopicIdIterator(
             Iterator<String> iterator,
-            TopicsImage image
+            TopicResolver resolver
         ) {
             this.iterator = Objects.requireNonNull(iterator);
-            this.image = Objects.requireNonNull(image);
+            this.resolver = Objects.requireNonNull(resolver);
         }
 
         @Override
@@ -85,9 +229,9 @@ public boolean hasNext() {
                     return false;
                 }
                 String next = iterator.next();
-                TopicImage topicImage = image.getTopic(next);
-                if (topicImage != null) {
-                    result = topicImage.id();
+                Uuid topicId = resolver.id(next);
+                if (topicId != null) {
+                    result = topicId;
                 }
             } while (result == null);
             next = result;
@@ -105,7 +249,7 @@ public Uuid next() {
 
     @Override
     public Iterator<Uuid> iterator() {
-        return new TopicIdIterator(topicNames.iterator(), image);
+        return new TopicIdIterator(topicNames.iterator(), resolver);
     }
 
     @Override
@@ -164,20 +308,20 @@ public boolean equals(Object o) {
         TopicIds uuids = (TopicIds) o;
 
         if (!Objects.equals(topicNames, uuids.topicNames)) return false;
-        return Objects.equals(image, uuids.image);
+        return Objects.equals(resolver, uuids.resolver);
     }
 
     @Override
     public int hashCode() {
         int result = topicNames.hashCode();
-        result = 31 * result + image.hashCode();
+        result = 31 * result + resolver.hashCode();
         return result;
     }
 
     @Override
     public String toString() {
         return "TopicIds(topicNames=" + topicNames +
-            ", image=" + image +
+            ", resolver=" + resolver +
             ')';
     }
 }

group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TargetAssignmentBuilderTest.java

@@ -161,6 +161,7 @@ public void prepareMemberAssignment(
 
         public TargetAssignmentBuilder.TargetAssignmentResult build() {
             TopicsImage topicsImage = topicsImageBuilder.build().topics();
+            TopicIds.TopicResolver topicResolver = new TopicIds.CachedTopicResolver(topicsImage);
             // Prepare expected member specs.
             Map<String, MemberSubscriptionAndAssignmentImpl> memberSubscriptions = new HashMap<>();
 
@@ -169,7 +170,7 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() {
                 memberSubscriptions.put(memberId, createMemberSubscriptionAndAssignment(
                     member,
                     targetAssignment.getOrDefault(memberId, Assignment.EMPTY),
-                    topicsImage
+                    topicResolver
                 ))
             );
 
@@ -192,7 +193,7 @@ public TargetAssignmentBuilder.TargetAssignmentResult build() {
                     memberSubscriptions.put(memberId, createMemberSubscriptionAndAssignment(
                         updatedMemberOrNull,
                         assignment,
-                        topicsImage
+                        topicResolver
                     ));
                 }
             });
@@ -263,6 +264,7 @@ public void testCreateMemberSubscriptionSpecImpl() {
             .addTopic(barTopicId, "bar", 5)
             .build()
             .topics();
+        TopicIds.TopicResolver topicResolver = new TopicIds.DefaultTopicResolver(topicsImage);
 
         ConsumerGroupMember member = new ConsumerGroupMember.Builder("member-id")
             .setSubscribedTopicNames(Arrays.asList("foo", "bar", "zar"))
@@ -278,7 +280,7 @@ public void testCreateMemberSubscriptionSpecImpl() {
         MemberSubscription subscriptionSpec = createMemberSubscriptionAndAssignment(
             member,
             assignment,
-            topicsImage
+            topicResolver
         );
 
         assertEquals(new MemberSubscriptionAndAssignmentImpl(

group-coordinator/src/test/java/org/apache/kafka/coordinator/group/modern/TopicIdsTest.java

@@ -41,7 +41,12 @@ public void testTopicNamesCannotBeNull() {
 
     @Test
     public void testTopicsImageCannotBeNull() {
-        assertThrows(NullPointerException.class, () -> new TopicIds(Collections.emptySet(), null));
+        assertThrows(NullPointerException.class, () -> new TopicIds(Collections.emptySet(), (TopicsImage) null));
+    }
+
+    @Test
+    public void testTopicResolverCannotBeNull() {
+        assertThrows(NullPointerException.class, () -> new TopicIds(Collections.emptySet(), (TopicIds.TopicResolver) null));
     }
 
     @Test

jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/AssignorBenchmarkUtils.java

@@ -159,13 +159,13 @@ public static TopicsImage createTopicsImage(Map<String, TopicMetadata> subscript
      *
      * @param members               The ConsumerGroupMembers.
      * @param subscriptionType      The group's subscription type.
-     * @param topicsImage           The TopicsImage to use.
+     * @param topicResolver         The TopicResolver to use.
      * @return The new GroupSpec.
      */
     public static GroupSpec createGroupSpec(
         Map<String, ConsumerGroupMember> members,
         SubscriptionType subscriptionType,
-        TopicsImage topicsImage
+        TopicIds.TopicResolver topicResolver
     ) {
         Map<String, MemberSubscriptionAndAssignmentImpl> memberSpecs = new HashMap<>();
 
@@ -177,7 +177,7 @@ public static GroupSpec createGroupSpec(
             memberSpecs.put(memberId, new MemberSubscriptionAndAssignmentImpl(
                 Optional.ofNullable(member.rackId()),
                 Optional.ofNullable(member.instanceId()),
-                new TopicIds(member.subscribedTopicNames(), topicsImage),
+                new TopicIds(member.subscribedTopicNames(), topicResolver),
                 new Assignment(member.assignedPartitions())
             ));
         }