feat: Centralize entity count tracking in EntityIdService

hapi/hedera-protobufs/block/stream/output/state_changes.proto

@@ -69,6 +69,7 @@ import "auxiliary/tss/tss_vote.proto";
 import "state/tss/tss_encryption_keys.proto";
 import "state/tss/tss_message_map_key.proto";
 import "state/tss/tss_vote_map_key.proto";
+import "state/entity/entity_counts.proto";
 
 /**
  * A set of state changes.
@@ -399,6 +400,11 @@ enum StateIdentifier {
      */
     STATE_ID_TSS_STATUS = 36;
 
+    /**
+     * A state identifier for the entity counts.
+     */
+    STATE_ID_ENTITY_COUNTS = 37;
+
     /**
      * A state identifier for the round receipts queue.
      */
@@ -608,6 +614,11 @@ message SingletonUpdateChange {
          * A change to the roster state singleton.
          */
         com.hedera.hapi.node.state.roster.RosterState roster_state_value = 13;
+
+        /**
+         * A change to the Entity counts singleton.
+         */
+        com.hedera.hapi.node.state.entity.EntityCounts entity_counts_value = 14;
     }
 }
 

hapi/hedera-protobufs/services/state/entity/entity_counts.proto

@@ -0,0 +1,89 @@
+syntax = "proto3";
+
+package com.hedera.hapi.node.state.entity;
+/*
+ * Copyright (C) 2025 Hedera Hashgraph, LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * This proto file contains primitive value messages.
+ * These are intended only for situations where the entire value to be stored in state is a single
+ * primitive.  These should never be used as components of another message; use the protobuf
+ * type instead.
+ */
+
+option java_package = "com.hederahashgraph.api.proto.java";
+// <<<pbj.java_package = "com.hedera.hapi.node.state.entity">>> This comment is special code for setting PBJ Compiler java package
+option java_multiple_files = true;
+
+/**
+ * Representation of a Hedera Entity Service entity counts in the network Merkle tree.
+ *
+ * This message is used to store the counts of various entities in the network.
+ */
+message EntityCounts {
+  /**
+   * The number of accounts in the network.
+   */
+  uint64 num_accounts = 1;
+  /**
+   * The number of aliases in the network.
+   */
+  uint64 num_aliases = 2;
+  /**
+   * The number of tokens in the network.
+   */
+  uint64 num_tokens = 3;
+  /**
+   * The number of token relationships in the network.
+   */
+  uint64 num_token_relations = 4;
+  /**
+   * The number of NFTs in the network.
+   */
+  uint64 num_nfts = 5;
+  /**
+   * The number of airdrops in the network.
+   */
+  uint64 num_airdrops = 6;
+  /**
+   * The number of staking infos in the network.
+   */
+  uint64 num_staking_infos = 7;
+  /**
+   * The number of topics in the network.
+   */
+  uint64 num_topics = 8;
+  /**
+   * The number of files in the network.
+   */
+  uint64 num_files = 9;
+  /**
+   * The number of nodes in the network.
+   */
+  uint64 num_nodes = 10;
+  /**
+   * The number of schedules in the network.
+   */
+  uint64 num_schedules = 11;
+  /**
+   * The number of contract storage slots in the network.
+   */
+  uint64 num_contract_storage_slots = 12;
+  /**
+   * The number of contract bytecodes in the network.
+   */
+  uint64 num_contract_bytecodes = 13;
+}

hapi/src/main/java/module-info.java

@@ -70,6 +70,7 @@
     exports com.hedera.hapi.block;
     exports com.hedera.hapi.services.auxiliary.tss.legacy;
     exports com.hedera.hapi.platform.event.legacy;
+    exports com.hedera.hapi.node.state.entity;
 
     requires transitive com.hedera.pbj.runtime;
     requires transitive com.google.common;

hedera-node/hedera-addressbook-service-impl/src/main/java/com/hedera/node/app/service/addressbook/impl/ReadableNodeStoreImpl.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2024-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -71,14 +71,6 @@ public Node get(final long nodeId) {
         return nodesState.get(EntityNumber.newBuilder().number(nodeId).build());
     }
 
-    /**
-     * Returns the number of topics in the state.
-     * @return the number of topics in the state
-     */
-    public long sizeOfState() {
-        return nodesState.size();
-    }
-
     protected <T extends ReadableKVState<EntityNumber, Node>> T nodesState() {
         return (T) nodesState;
     }

hedera-node/hedera-addressbook-service-impl/src/main/java/com/hedera/node/app/service/addressbook/impl/WritableNodeStore.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2024-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,10 +20,6 @@
 
 import com.hedera.hapi.node.state.addressbook.Node;
 import com.hedera.hapi.node.state.common.EntityNumber;
-import com.hedera.node.app.spi.metrics.StoreMetricsService;
-import com.hedera.node.app.spi.metrics.StoreMetricsService.StoreType;
-import com.hedera.node.config.data.NodesConfig;
-import com.swirlds.config.api.Configuration;
 import com.swirlds.state.spi.WritableKVState;
 import com.swirlds.state.spi.WritableStates;
 import edu.umd.cs.findbugs.annotations.NonNull;
@@ -41,18 +37,9 @@ public class WritableNodeStore extends ReadableNodeStoreImpl {
      * Create a new {@link WritableNodeStore} instance.
      *
      * @param states The state to use.
-     * @param configuration The configuration used to read the maximum capacity.
-     * @param storeMetricsService Service that provides utilization metrics.
      */
-    public WritableNodeStore(
-            @NonNull final WritableStates states,
-            @NonNull final Configuration configuration,
-            @NonNull final StoreMetricsService storeMetricsService) {
+    public WritableNodeStore(@NonNull final WritableStates states) {
         super(states);
-
-        final long maxCapacity = configuration.getConfigData(NodesConfig.class).maxNumber();
-        final var storeMetrics = storeMetricsService.get(StoreType.NODE, maxCapacity);
-        nodesState().setMetrics(storeMetrics);
     }
 
     @Override
@@ -81,15 +68,6 @@ public Node getForModify(final long nodeId) {
                 .getForModify(EntityNumber.newBuilder().number(nodeId).build());
     }
 
-    /**
-     * Returns the number of nodes in the state.
-     * @return the number of nodes in the state
-     */
-    @Override
-    public long sizeOfState() {
-        return nodesState().size();
-    }
-
     /**
      * Returns the set of nodes modified in existing state.
      * @return the set of nodes modified in existing state

hedera-node/hedera-addressbook-service-impl/src/main/java/com/hedera/node/app/service/addressbook/impl/handlers/NodeCreateHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2024-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -41,6 +41,7 @@
 import com.hedera.node.app.service.token.ReadableAccountStore;
 import com.hedera.node.app.spi.fees.FeeContext;
 import com.hedera.node.app.spi.fees.Fees;
+import com.hedera.node.app.spi.ids.ReadableEntityIdStore;
 import com.hedera.node.app.spi.workflows.HandleContext;
 import com.hedera.node.app.spi.workflows.PreCheckException;
 import com.hedera.node.app.spi.workflows.PreHandleContext;
@@ -100,9 +101,10 @@ public void handle(@NonNull final HandleContext handleContext) {
         final var storeFactory = handleContext.storeFactory();
         final var nodeStore = storeFactory.writableStore(WritableNodeStore.class);
         final var accountStore = storeFactory.readableStore(ReadableAccountStore.class);
+        final var entityIdStore = storeFactory.readableStore(ReadableEntityIdStore.class);
         final var accountId = op.accountIdOrElse(AccountID.DEFAULT);
 
-        validateFalse(nodeStore.sizeOfState() >= nodeConfig.maxNumber(), MAX_NODES_CREATED);
+        validateFalse(entityIdStore.numNodes() >= nodeConfig.maxNumber(), MAX_NODES_CREATED);
         validateTrue(accountStore.contains(accountId), INVALID_NODE_ACCOUNT_ID);
         addressBookValidator.validateDescription(op.description(), nodeConfig);
         addressBookValidator.validateGossipEndpoint(op.gossipEndpoint(), nodeConfig);

hedera-node/hedera-addressbook-service-impl/src/test/java/com/hedera/node/app/service/addressbook/impl/test/ReadableNodeStoreImplTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -87,12 +87,6 @@ void nullArgsFail() {
         assertThrows(NullPointerException.class, () -> new ReadableNodeStoreImpl(null));
     }
 
-    @Test
-    void getSizeOfState() {
-        final var store = new ReadableNodeStoreImpl(readableStates);
-        assertEquals(readableStates.get(NODES_KEY).size(), store.sizeOfState());
-    }
-
     @Test
     void keysWorks() {
         final var stateBuilder = emptyReadableNodeStateBuilder();

hedera-node/hedera-addressbook-service-impl/src/test/java/com/hedera/node/app/service/addressbook/impl/test/WritableNodeStoreTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -45,16 +45,15 @@ class WritableNodeStoreTest extends AddressBookTestBase {
 
     @Test
     void throwsIfNullValuesAsArgs() {
-        assertThrows(NullPointerException.class, () -> new WritableNodeStore(null, CONFIGURATION, storeMetricsService));
-        assertThrows(
-                NullPointerException.class, () -> new WritableNodeStore(writableStates, null, storeMetricsService));
-        assertThrows(NullPointerException.class, () -> new WritableNodeStore(writableStates, CONFIGURATION, null));
+        assertThrows(NullPointerException.class, () -> new WritableNodeStore(null));
+        assertThrows(NullPointerException.class, () -> new WritableNodeStore(writableStates));
+        assertThrows(NullPointerException.class, () -> new WritableNodeStore(writableStates));
         assertThrows(NullPointerException.class, () -> writableStore.put(null));
     }
 
     @Test
     void constructorCreatesNodeState() {
-        final var store = new WritableNodeStore(writableStates, CONFIGURATION, storeMetricsService);
+        final var store = new WritableNodeStore(writableStates);
         assertNotNull(store);
     }
 

hedera-node/hedera-addressbook-service-impl/src/test/java/com/hedera/node/app/service/addressbook/impl/test/handlers/AddressBookTestBase.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -178,7 +178,7 @@ protected void refreshStoresWithCurrentNodeInReadable() {
         given(readableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(readableNodeState);
         given(writableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(writableNodeState);
         readableStore = new ReadableNodeStoreImpl(readableStates);
-        writableStore = new WritableNodeStore(writableStates, DEFAULT_CONFIG, storeMetricsService);
+        writableStore = new WritableNodeStore(writableStates);
     }
 
     protected void refreshStoresWithCurrentNodeInBothReadableAndWritable() {
@@ -188,21 +188,21 @@ protected void refreshStoresWithCurrentNodeInBothReadableAndWritable() {
         given(writableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(writableNodeState);
         readableStore = new ReadableNodeStoreImpl(readableStates);
         final var configuration = HederaTestConfigBuilder.createConfig();
-        writableStore = new WritableNodeStore(writableStates, configuration, storeMetricsService);
+        writableStore = new WritableNodeStore(writableStates);
     }
 
     protected void refreshStoresWithCurrentNodeInWritable() {
         writableNodeState = writableNodeStateWithOneKey();
         given(writableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(writableNodeState);
         final var configuration = HederaTestConfigBuilder.createConfig();
-        writableStore = new WritableNodeStore(writableStates, configuration, storeMetricsService);
+        writableStore = new WritableNodeStore(writableStates);
     }
 
     protected void refreshStoresWithMoreNodeInWritable() {
         writableNodeState = writableNodeStateWithMoreKeys();
         given(writableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(writableNodeState);
         final var configuration = HederaTestConfigBuilder.createConfig();
-        writableStore = new WritableNodeStore(writableStates, configuration, storeMetricsService);
+        writableStore = new WritableNodeStore(writableStates);
     }
 
     @NonNull

hedera-node/hedera-addressbook-service-impl/src/test/java/com/hedera/node/app/service/addressbook/impl/test/handlers/NodeCreateHandlerTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -232,7 +232,7 @@ void failsWhenMaxNodesExceeds() {
         given(handleContext.storeFactory()).willReturn(storeFactory);
         given(storeFactory.writableStore(WritableNodeStore.class)).willReturn(writableStore);
 
-        assertEquals(1, writableStore.sizeOfState());
+        //        assertEquals(1, writableStore.sizeOfState());
         final var msg = assertThrows(HandleException.class, () -> subject.handle(handleContext));
         assertEquals(ResponseCodeEnum.MAX_NODES_CREATED, msg.getStatus());
         assertEquals(0, writableStore.modifiedNodes().size());

hedera-node/hedera-addressbook-service-impl/src/test/java/com/hedera/node/app/service/addressbook/impl/test/handlers/NodeDeleteHandlerTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -97,7 +97,7 @@ void setUp() {
         writableNodeState = writableNodeStateWithOneKey();
         given(writableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(writableNodeState);
         testConfig = HederaTestConfigBuilder.createConfig();
-        writableStore = new WritableNodeStore(writableStates, testConfig, storeMetricsService);
+        writableStore = new WritableNodeStore(writableStates);
         lenient().when(handleContext.configuration()).thenReturn(testConfig);
     }
 
@@ -150,7 +150,7 @@ void fileDoesntExist() {
         given(handleContext.storeFactory()).willReturn(storeFactory);
         writableNodeState = emptyWritableNodeState();
         given(writableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(writableNodeState);
-        writableStore = new WritableNodeStore(writableStates, testConfig, storeMetricsService);
+        writableStore = new WritableNodeStore(writableStates);
         given(storeFactory.writableStore(WritableNodeStore.class)).willReturn(writableStore);
 
         given(handleContext.body())
@@ -171,7 +171,7 @@ void NodeIsNull() {
         given(handleContext.storeFactory()).willReturn(storeFactory);
         writableNodeState = writableNodeStateWithOneKey();
         given(writableStates.<EntityNumber, Node>get(NODES_KEY)).willReturn(writableNodeState);
-        writableStore = new WritableNodeStore(writableStates, testConfig, storeMetricsService);
+        writableStore = new WritableNodeStore(writableStates);
         given(storeFactory.writableStore(WritableNodeStore.class)).willReturn(writableStore);
 
         given(handleContext.body())

hedera-node/hedera-addressbook-service/src/main/java/com/hedera/node/app/service/addressbook/AddressBookHelper.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 Hedera Hashgraph, LLC
+ * Copyright (C) 2024-2025 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -61,7 +61,7 @@ private AddressBookHelper() {
     public static long getNextNodeID(@NonNull final ReadableNodeStore nodeStore) {
         requireNonNull(nodeStore);
         final long maxNodeId = StreamSupport.stream(
-                        Spliterators.spliterator(nodeStore.keys(), nodeStore.sizeOfState(), DISTINCT), false)
+                        Spliterators.spliteratorUnknownSize(nodeStore.keys(), DISTINCT), false)
                 .mapToLong(EntityNumber::number)
                 .max()
                 .orElse(-1L);