adding network size query (#471)

JackalLabs · Aug 19, 2024 · 3b31b0f · 3b31b0f
2 parents 40d9585 + 312a976
commit 3b31b0f
Show file tree

Hide file tree

Showing 5 changed files with 585 additions and 149 deletions.
diff --git a/proto/canine_chain/storage/query.proto b/proto/canine_chain/storage/query.proto
@@ -181,6 +181,12 @@ service Query {
         "/jackal/canine-chain/storage/storage_stats";
   }
 
+  // Queries protocol storage space used and purchased.
+  rpc NetworkSize(QueryNetworkSize) returns (QueryNetworkSizeResponse) {
+    option (google.api.http).get =
+        "/jackal/canine-chain/storage/network_size";
+  }
+
   // Queries protocol storage space used and purchased.
   rpc Gauges(QueryAllGauges) returns (QueryAllGaugesResponse) {
     option (google.api.http).get =
@@ -440,3 +446,9 @@ message QueryStorageStatsResponse {
   uint64 uniqueUsers = 5;
   map<int64, int64> users_by_plan= 6;
 }
+
+message QueryNetworkSize {}
+
+message QueryNetworkSizeResponse {
+  uint64 size = 1;
+}
diff --git a/x/storage/keeper/files.go b/x/storage/keeper/files.go
@@ -138,6 +138,33 @@ func (k Keeper) IterateFilesByMerkle(ctx sdk.Context, reverse bool, fn func(key
 	}
 }
 
+// IterateAndParseFilesByMerkle iterates through every file and parses them for you
+func (k Keeper) IterateAndParseFilesByMerkle(ctx sdk.Context, reverse bool, fn func(key []byte, val types.UnifiedFile) bool) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefix(types.FilePrimaryKeyPrefix))
+
+	var iterator storetypes.Iterator
+	if reverse {
+		iterator = sdk.KVStoreReversePrefixIterator(store, []byte{})
+	} else {
+		iterator = sdk.KVStorePrefixIterator(store, []byte{})
+	}
+
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		val := iterator.Value()
+		var file types.UnifiedFile
+		if err := k.cdc.Unmarshal(val, &file); err != nil {
+			return
+		}
+
+		b := fn(iterator.Key(), file)
+		if b {
+			return
+		}
+	}
+}
+
 // GetAllFilesWithMerkle returns all Files that start with a specific merkle
 func (k Keeper) GetAllFilesWithMerkle(ctx sdk.Context, merkle []byte) (list []types.UnifiedFile) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.FilesMerklePrefix(merkle))

diff --git a/x/storage/keeper/grpc_query_storage_stats.go b/x/storage/keeper/grpc_query_storage_stats.go
@@ -9,6 +9,24 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+func (k Keeper) NetworkSize(c context.Context, req *types.QueryNetworkSize) (*types.QueryNetworkSizeResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	}
+
+	ctx := sdk.UnwrapSDKContext(c)
+
+	var s uint64
+
+	k.IterateAndParseFilesByMerkle(ctx, false, func(_ []byte, val types.UnifiedFile) bool {
+		s += uint64(val.FileSize * val.MaxProofs)
+
+		return false
+	})
+
+	return &types.QueryNetworkSizeResponse{Size_: s}, nil
+}
+
 func (k Keeper) StorageStats(c context.Context, req *types.QueryStorageStats) (*types.QueryStorageStatsResponse, error) {
 	if req == nil {
 		return nil, status.Error(codes.InvalidArgument, "invalid request")