feat: add getMemoryStats endpoint to registry

backend/main/DownloadLog.mo

@@ -16,6 +16,7 @@ import DateComponents "mo:datetime/Components";
 
 import Types "./types";
 import PackageUtils "./utils/package-utils";
+import MemoryStats "./MemoryStats";
 
 module {
   public type DownloadsSnapshot = Types.DownloadsSnapshot;
@@ -419,6 +420,33 @@ module {
       );
     };
 
+    public func getMemoryStats() : {
+      downloadsByPackageName : MemoryStats.StructureStats;
+      downloadsByPackageId : MemoryStats.StructureStats;
+      dailySnapshots : MemoryStats.StructureStats;
+      weeklySnapshots : MemoryStats.StructureStats;
+      dailySnapshotsByPackageName : MemoryStats.StructureStats;
+      dailySnapshotsByPackageId : MemoryStats.StructureStats;
+      weeklySnapshotsByPackageName : MemoryStats.StructureStats;
+      weeklySnapshotsByPackageId : MemoryStats.StructureStats;
+      dailyTempRecords : MemoryStats.StructureStats;
+      weeklyTempRecords : MemoryStats.StructureStats;
+    } {
+      let serializeSnapshot = func(v : DownloadsSnapshot) : Blob = to_candid (v);
+      {
+        downloadsByPackageName = MemoryStats.statsForMap(downloadsByPackageName, func(k : Text.Text, v : Nat) : Blob = to_candid ((k, v)));
+        downloadsByPackageId = MemoryStats.statsForMap(downloadsByPackageId, func(k : Text.Text, v : Nat) : Blob = to_candid ((k, v)));
+        dailySnapshots = MemoryStats.statsForBuffer(dailySnapshots, serializeSnapshot);
+        weeklySnapshots = MemoryStats.statsForBuffer(weeklySnapshots, serializeSnapshot);
+        dailySnapshotsByPackageName = MemoryStats.statsForMapOfBuffers(dailySnapshotsByPackageName, serializeSnapshot);
+        dailySnapshotsByPackageId = MemoryStats.statsForMapOfBuffers(dailySnapshotsByPackageId, serializeSnapshot);
+        weeklySnapshotsByPackageName = MemoryStats.statsForMapOfBuffers(weeklySnapshotsByPackageName, serializeSnapshot);
+        weeklySnapshotsByPackageId = MemoryStats.statsForMapOfBuffers(weeklySnapshotsByPackageId, serializeSnapshot);
+        dailyTempRecords = MemoryStats.statsForBuffer(dailyTempRecords, func(v : Record) : Blob = to_candid (v));
+        weeklyTempRecords = MemoryStats.statsForBuffer(weeklyTempRecords, func(v : Record) : Blob = to_candid (v));
+      };
+    };
+
     public func setTimers<system>() {
       cancelTimers();
       timerId := Timer.recurringTimer<system>(

backend/main/MemoryStats.mo

@@ -0,0 +1,155 @@
+import Nat "mo:base/Nat";
+import Buffer "mo:base/Buffer";
+import TrieMap "mo:base/TrieMap";
+import Iter "mo:base/Iter";
+
+module {
+  public type StructureStats = {
+    count : Nat;
+    bytes : Nat;
+  };
+
+  public type MemoryStats = {
+    rtsHeapSize : Nat;
+    rtsMemorySize : Nat;
+
+    packageVersions : StructureStats;
+    packageConfigs : StructureStats;
+    highestConfigs : StructureStats;
+    packagePublications : StructureStats;
+    ownersByPackage : StructureStats;
+    maintainersByPackage : StructureStats;
+    fileIdsByPackage : StructureStats;
+    hashByFileId : StructureStats;
+    packageFileStats : StructureStats;
+    packageTestStats : StructureStats;
+    packageBenchmarks : StructureStats;
+    packageNotes : StructureStats;
+    packageDocsCoverage : StructureStats;
+
+    downloadsByPackageName : StructureStats;
+    downloadsByPackageId : StructureStats;
+    dailySnapshots : StructureStats;
+    weeklySnapshots : StructureStats;
+    dailySnapshotsByPackageName : StructureStats;
+    dailySnapshotsByPackageId : StructureStats;
+    weeklySnapshotsByPackageName : StructureStats;
+    weeklySnapshotsByPackageId : StructureStats;
+    dailyTempRecords : StructureStats;
+    weeklyTempRecords : StructureStats;
+
+    storages : StructureStats;
+    storageByFileId : StructureStats;
+
+    users : StructureStats;
+    names : StructureStats;
+  };
+
+  public let SAMPLE_SIZE : Nat = 1_000;
+
+  // Serializes a systematic sample of up to SAMPLE_SIZE elements from an
+  // iterator and extrapolates the total byte count, keeping to_candid
+  // allocations bounded.
+  public func sampleIterBytes<V>(
+    iter : Iter.Iter<V>,
+    total : Nat,
+    serialize : V -> Blob,
+  ) : Nat {
+    if (total == 0) return 0;
+    let stride = if (total <= SAMPLE_SIZE) 1 else (total + SAMPLE_SIZE - 1 : Nat) / SAMPLE_SIZE;
+    var sum = 0;
+    var i = 0;
+    var sampled = 0;
+    label sampleLoop for (v in iter) {
+      if (sampled >= SAMPLE_SIZE) break sampleLoop;
+      if (i % stride == 0) {
+        sum += serialize(v).size();
+        sampled += 1;
+      };
+      i += 1;
+    };
+    if (sampled == 0) 0 else sum * total / sampled;
+  };
+
+  // Two-level sampling for TrieMaps whose values are Buffers.
+  // Budgets to_candid calls across both keys and their inner elements.
+  public func sampleMapOfBuffersBytes<K, V>(
+    map : TrieMap.TrieMap<K, Buffer.Buffer<V>>,
+    serialize : V -> Blob,
+  ) : Nat {
+    let total = map.size();
+    if (total == 0) return 0;
+    let numSampledKeys = Nat.min(total, SAMPLE_SIZE);
+    let perKeyBudget = Nat.max(1, SAMPLE_SIZE / numSampledKeys);
+    let stride = if (total <= SAMPLE_SIZE) 1 else (total + SAMPLE_SIZE - 1 : Nat) / SAMPLE_SIZE;
+    var sum = 0;
+    var i = 0;
+    var sampled = 0;
+    label sampleLoop for ((_, buf) in map.entries()) {
+      if (sampled >= SAMPLE_SIZE) break sampleLoop;
+      if (i % stride == 0) {
+        let bufTotal = buf.size();
+        if (bufTotal > 0) {
+          let bufStride = if (bufTotal <= perKeyBudget) 1 else (bufTotal + perKeyBudget - 1 : Nat) / perKeyBudget;
+          var j = 0;
+          var bufSum = 0;
+          var bufSampled = 0;
+          while (j < bufTotal and bufSampled < perKeyBudget) {
+            bufSum += serialize(buf.get(j)).size();
+            bufSampled += 1;
+            j += bufStride;
+          };
+          sum += bufSum * bufTotal / bufSampled;
+        };
+        sampled += 1;
+      };
+      i += 1;
+    };
+    if (sampled == 0) return 0;
+    sum * total / sampled;
+  };
+
+  // Convenience: returns { count; bytes } for a TrieMap.
+  public func statsForMap<K, V>(
+    map : TrieMap.TrieMap<K, V>,
+    serialize : (K, V) -> Blob,
+  ) : StructureStats {
+    let total = map.size();
+    {
+      count = total;
+      bytes = sampleIterBytes(
+        map.entries(),
+        total,
+        func((k, v) : (K, V)) : Blob = serialize(k, v),
+      );
+    };
+  };
+
+  // Convenience: returns { count; bytes } for a Buffer.
+  public func statsForBuffer<V>(
+    buf : Buffer.Buffer<V>,
+    serialize : V -> Blob,
+  ) : StructureStats {
+    {
+      count = buf.size();
+      bytes = sampleIterBytes(buf.vals(), buf.size(), serialize);
+    };
+  };
+
+  // Convenience: returns { count; bytes } for an arbitrary Iter + known size.
+  public func statsForIter<V>(
+    iter : Iter.Iter<V>,
+    total : Nat,
+    serialize : V -> Blob,
+  ) : StructureStats {
+    { count = total; bytes = sampleIterBytes(iter, total, serialize) };
+  };
+
+  // Convenience: returns { count; bytes } for a TrieMap<K, Buffer<V>>.
+  public func statsForMapOfBuffers<K, V>(
+    map : TrieMap.TrieMap<K, Buffer.Buffer<V>>,
+    serialize : V -> Blob,
+  ) : StructureStats {
+    { count = map.size(); bytes = sampleMapOfBuffersBytes(map, serialize) };
+  };
+};

backend/main/Users.mo

@@ -12,6 +12,7 @@ import Set "mo:map/Set";
 
 import Types "./types";
 import { isLetter; isLowerCaseLetter } "./utils/is-letter";
+import MemoryStats "./MemoryStats";
 
 module {
   public type Stable = ?{
@@ -25,6 +26,16 @@ module {
     var _users = TrieMap.TrieMap<Principal, Types.User>(Principal.equal, Principal.hash);
     var _names = Set.new<Text>();
 
+    public func getMemoryStats() : {
+      users : MemoryStats.StructureStats;
+      names : MemoryStats.StructureStats;
+    } {
+      {
+        users = MemoryStats.statsForMap(_users, func(k : Principal, v : Types.User) : Blob = to_candid ((k, v)));
+        names = MemoryStats.statsForIter(Set.keys(_names), Set.size(_names), func(k : Text) : Blob = to_candid (k));
+      };
+    };
+
     public func toStable() : Stable {
       ?#v2({
         users = Iter.toArray(_users.entries());

backend/main/main-canister.mo

@@ -13,6 +13,7 @@ import Option "mo:base/Option";
 import Blob "mo:base/Blob";
 import TelegramBot "mo:telegram-bot";
 
+import { rts_heap_size; rts_memory_size } "mo:prim";
 import IC "mo:ic";
 import { DAY } "mo:time-consts";
 import Backup "mo:backup";
@@ -29,6 +30,7 @@ import Storage "../storage/storage-canister";
 import Users "./Users";
 import Badges "./badges";
 
+import MemoryStats "./MemoryStats";
 import Registry "./registry/Registry";
 import PackagePublisher "./PackagePublisher";
 import { searchInRegistry } "./registry/searchInRegistry";
@@ -579,6 +581,34 @@ actor class Main() = this {
     storageManager.getStoragesStats();
   };
 
+  public query ({ caller }) func getMemoryStats() : async MemoryStats.MemoryStats {
+    assert (Utils.isAdmin(caller));
+
+    let dlStats = downloadLog.getMemoryStats();
+    let smStats = storageManager.getMemoryStats();
+    let uStats = users.getMemoryStats();
+
+    {
+      dlStats and smStats and uStats with
+      rtsHeapSize = rts_heap_size();
+      rtsMemorySize = rts_memory_size();
+
+      packageVersions = MemoryStats.statsForMap(packageVersions, func(k : PackageName, v : [PackageVersion]) : Blob = to_candid ((k, v)));
+      packageConfigs = MemoryStats.statsForMap(packageConfigs, func(k : PackageId, v : PackageConfigV3) : Blob = to_candid ((k, v)));
+      highestConfigs = MemoryStats.statsForMap(highestConfigs, func(k : PackageName, v : PackageConfigV3) : Blob = to_candid ((k, v)));
+      packagePublications = MemoryStats.statsForMap(packagePublications, func(k : PackageId, v : PackagePublication) : Blob = to_candid ((k, v)));
+      ownersByPackage = MemoryStats.statsForMap(ownersByPackage, func(k : PackageName, v : [Principal]) : Blob = to_candid ((k, v)));
+      maintainersByPackage = MemoryStats.statsForMap(maintainersByPackage, func(k : PackageName, v : [Principal]) : Blob = to_candid ((k, v)));
+      fileIdsByPackage = MemoryStats.statsForMap(fileIdsByPackage, func(k : PackageId, v : [FileId]) : Blob = to_candid ((k, v)));
+      hashByFileId = MemoryStats.statsForMap(hashByFileId, func(k : FileId, v : Blob) : Blob = to_candid ((k, v)));
+      packageFileStats = MemoryStats.statsForMap(packageFileStats, func(k : PackageId, v : PackageFileStats) : Blob = to_candid ((k, v)));
+      packageTestStats = MemoryStats.statsForMap(packageTestStats, func(k : PackageId, v : TestStats) : Blob = to_candid ((k, v)));
+      packageBenchmarks = MemoryStats.statsForMap(packageBenchmarks, func(k : PackageId, v : Benchmarks) : Blob = to_candid ((k, v)));
+      packageNotes = MemoryStats.statsForMap(packageNotes, func(k : PackageId, v : Text) : Blob = to_candid ((k, v)));
+      packageDocsCoverage = MemoryStats.statsForMap(packageDocsCoverage, func(k : PackageId, v : Float) : Blob = to_candid ((k, v)));
+    };
+  };
+
   // USERS
   public query func getUser(userId : Principal) : async ?User {
     users.getUserOpt(userId);

backend/storage/storage-manager.mo

@@ -11,6 +11,7 @@ import { ic } "mo:ic";
 
 import Types "./types";
 import Storage "./storage-canister";
+import MemoryStats "../main/MemoryStats";
 
 module {
   public type StorageId = Principal;
@@ -172,6 +173,16 @@ module {
       );
     };
 
+    public func getMemoryStats() : {
+      storages : MemoryStats.StructureStats;
+      storageByFileId : MemoryStats.StructureStats;
+    } {
+      {
+        storages = MemoryStats.statsForMap(storages, func(k : StorageId, v : StorageStats) : Blob = to_candid ((k, v)));
+        storageByFileId = MemoryStats.statsForMap(storageByFileId, func(k : FileId, v : StorageId) : Blob = to_candid ((k, v)));
+      };
+    };
+
     public func toStable() : Stable {
       return ?{
         storages = Iter.toArray(storages.entries());

cli/declarations/main/main.did

@@ -42,6 +42,11 @@ type StreamingCallbackResponse =
  };
 type StreamingCallback = func (StreamingToken) ->
                           (opt StreamingCallbackResponse) query;
+type StructureStats = 
+ record {
+   bytes: nat;
+   count: nat;
+ };
 type StorageStats = 
  record {
    cyclesBalance: nat;
@@ -272,6 +277,38 @@ type PackageChanges =
    prevDocsCoverage: float64;
    tests: TestsChanges;
  };
+type MemoryStats = 
+ record {
+   dailySnapshots: StructureStats;
+   dailySnapshotsByPackageId: StructureStats;
+   dailySnapshotsByPackageName: StructureStats;
+   dailyTempRecords: StructureStats;
+   downloadsByPackageId: StructureStats;
+   downloadsByPackageName: StructureStats;
+   fileIdsByPackage: StructureStats;
+   hashByFileId: StructureStats;
+   highestConfigs: StructureStats;
+   maintainersByPackage: StructureStats;
+   ownersByPackage: StructureStats;
+   packageBenchmarks: StructureStats;
+   packageConfigs: StructureStats;
+   packageDocsCoverage: StructureStats;
+   packageFileStats: StructureStats;
+   packageNotes: StructureStats;
+   packagePublications: StructureStats;
+   packageTestStats: StructureStats;
+   packageVersions: StructureStats;
+   rtsHeapSize: nat;
+   rtsMemorySize: nat;
+   names: StructureStats;
+   storageByFileId: StructureStats;
+   storages: StructureStats;
+   users: StructureStats;
+   weeklySnapshots: StructureStats;
+   weeklySnapshotsByPackageId: StructureStats;
+   weeklySnapshotsByPackageName: StructureStats;
+   weeklyTempRecords: StructureStats;
+ };
 type Main = 
  service {
    addMaintainer: (packageName: PackageName, newMaintainer: principal) ->
@@ -311,6 +348,7 @@ type Main =
           SemverPart;
         }) -> (Result_6) query;
    getHighestVersion: (name: PackageName) -> (Result_5) query;
+   getMemoryStats: () -> (MemoryStats) query;
    getMostDownloadedPackages: () -> (vec PackageSummary) query;
    getMostDownloadedPackagesIn7Days: () -> (vec PackageSummary) query;
    getNewPackages: () -> (vec PackageSummary) query;