From f80d42d6f334cf461b80d5f3f4a23be6560e1d41 Mon Sep 17 00:00:00 2001
From: Scott Fauerbach <scottfauerbach@gmail.com>
Date: Wed, 10 Jul 2024 15:51:26 -0400
Subject: [PATCH] Stream state subjects as a map (#1177)

---
 .../java/io/nats/client/api/StreamState.java  | 20 +++++++++++++++++--
 .../client/impl/JetStreamManagementTests.java |  3 +++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/main/java/io/nats/client/api/StreamState.java b/src/main/java/io/nats/client/api/StreamState.java
index 37027e549..81e38dee3 100644
--- a/src/main/java/io/nats/client/api/StreamState.java
+++ b/src/main/java/io/nats/client/api/StreamState.java
@@ -16,7 +16,9 @@
 import io.nats.client.support.JsonValue;
 
 import java.time.ZonedDateTime;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static io.nats.client.support.ApiConstants.*;
 import static io.nats.client.support.JsonValueUtils.*;
@@ -34,6 +36,7 @@ public class StreamState {
     private final List<Subject> subjects;
     private final List<Long> deletedStreamSequences;
     private final LostStreamData lostStreamData;
+    private final Map<String, Long> subjectMap;
 
     StreamState(JsonValue vStreamState) {
         msgs = readLong(vStreamState, MESSAGES, 0);
@@ -48,6 +51,11 @@ public class StreamState {
         subjects = Subject.listOf(readValue(vStreamState, SUBJECTS));
         deletedStreamSequences = readLongList(vStreamState, DELETED);
         lostStreamData = LostStreamData.optionalInstance(readValue(vStreamState, LOST));
+
+        subjectMap = new HashMap<>();
+        for (Subject s : subjects) {
+            subjectMap.put(s.getName(), s.getCount());
+        }
     }
 
     /**
@@ -122,14 +130,22 @@ public long getSubjectCount() {
     }
 
     /**
-     * Get a list of the Subject objects. May be null if the Stream Info request did not ask for subjects
-     * or if there are no subjects.
+     * Get a list of the Subject objects. May be empty, for instance
+     * if the Stream Info request did not ask for subjects or if there are no subjects.
      * @return the list of subjects
      */
     public List<Subject> getSubjects() {
         return subjects;
     }
 
+    /**
+     * Get a map of subjects instead of a list of Subject objects.
+     * @return the map
+     */
+    public Map<String, Long> getSubjectMap() {
+        return subjectMap;
+    }
+
     /**
      * Gets the count of deleted messages
      *
diff --git a/src/test/java/io/nats/client/impl/JetStreamManagementTests.java b/src/test/java/io/nats/client/impl/JetStreamManagementTests.java
index 4e8e8819c..b929a6bc0 100644
--- a/src/test/java/io/nats/client/impl/JetStreamManagementTests.java
+++ b/src/test/java/io/nats/client/impl/JetStreamManagementTests.java
@@ -403,6 +403,7 @@ public void testGetStreamInfo() throws Exception {
             assertEquals(0, si.getStreamState().getSubjects().size());
             assertEquals(0, si.getStreamState().getDeletedCount());
             assertEquals(0, si.getStreamState().getDeleted().size());
+            assertTrue(si.getStreamState().getSubjectMap().isEmpty());
 
             if (nc.getServerInfo().isOlderThanVersion("2.10")) {
                 assertNull(si.getTimestamp());
@@ -427,6 +428,7 @@ public void testGetStreamInfo() throws Exception {
             assertEquals(0, si.getStreamState().getSubjects().size());
             assertEquals(5, si.getStreamState().getDeletedCount());
             assertEquals(0, si.getStreamState().getDeleted().size());
+            assertTrue(si.getStreamState().getSubjectMap().isEmpty());
 
             si = jsm.getStreamInfo(stream, StreamInfoOptions.builder().allSubjects().deletedDetails().build());
             assertEquals(stream, si.getConfiguration().getName());
@@ -448,6 +450,7 @@ public void testGetStreamInfo() throws Exception {
             Subject sf = map.get(subjectIx5 + ".bar");
             assertNotNull(sf);
             assertEquals(6, sf.getCount());
+            assertEquals(6, si.getStreamState().getSubjectMap().size());
 
             for (PublishAck pa : packs) {
                 assertTrue(si.getStreamState().getDeleted().contains(pa.getSeqno()));