Use new ReconciliationAnnounceEntries/-TerminatePayload format

xash · Nov 13, 2024 · 94eed96 · 94eed96
1 parent 2c6df3e
commit 94eed96
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 22 deletions.
diff --git a/src/wgps/reconciliation/announcer.ts b/src/wgps/reconciliation/announcer.ts
@@ -34,7 +34,7 @@ type AnnouncementPack<
   // Then send a ReconciliationAnnounceEntries
   announcement: {
     range: Range3d<SubspaceId>;
-    count: number;
+    isEmpty: boolean;
     wantResponse: boolean;
     senderHandle: bigint;
     receiverHandle: bigint;
@@ -176,7 +176,7 @@ export class Announcer<
     this.announcementPackQueue.push({
       staticTokenBinds,
       announcement: {
-        count: entries.length,
+        isEmpty: entries.length === 0,
         range: announcement.range,
         receiverHandle: announcement.receiverHandle,
         senderHandle: announcement.senderHandle,

diff --git a/src/wgps/reconciliation/reconcile_msg_tracker.ts b/src/wgps/reconciliation/reconcile_msg_tracker.ts
@@ -8,6 +8,7 @@ import type {
   MsgReconciliationAnnounceEntries,
   MsgReconciliationSendEntry,
   MsgReconciliationSendFingerprint,
+  MsgReconciliationTerminatePayload,
   ReconciliationPrivy,
 } from "../types.ts";
 
@@ -37,7 +38,7 @@ export class ReconcileMsgTracker<
 
   private announcedRange: Range3d<SubspaceId>;
   private announcedNamespace: NamespaceId;
-  private announcedEntriesRemaining = 0n;
+  private announcedEntriesExpecting = false;
 
   private handleToNamespaceId: (aoiHandle: bigint) => NamespaceId;
 
@@ -76,7 +77,7 @@ export class ReconcileMsgTracker<
 
     this.announcedRange = msg.range;
     this.announcedNamespace = this.handleToNamespaceId(msg.receiverHandle);
-    this.announcedEntriesRemaining = msg.count;
+    this.announcedEntriesExpecting = !msg.isEmpty;
   }
 
   onSendEntry(
@@ -90,25 +91,22 @@ export class ReconcileMsgTracker<
     this.prevEntry = msg.entry.entry;
     this.prevToken = msg.staticTokenHandle;
 
-    this.announcedEntriesRemaining -= 1n;
-
     this.isAwaitingTermination = true;
   }
 
-  onTerminatePayload() {
+  onTerminatePayload(msg: MsgReconciliationTerminatePayload) {
     this.isAwaitingTermination = false;
+    if (msg.isFinal) {
+      this.announcedEntriesExpecting = false;
+    }
   }
 
   isExpectingPayloadOrTermination() {
     return this.isAwaitingTermination;
   }
 
   isExpectingReconciliationSendEntry() {
-    if (this.announcedEntriesRemaining > 0n) {
-      return true;
-    }
-
-    return false;
+    return this.announcedEntriesExpecting;
   }
 
   getPrivy(): ReconciliationPrivy<NamespaceId, SubspaceId, PayloadDigest> {

diff --git a/src/wgps/reconciliation/reconciler.ts b/src/wgps/reconciliation/reconciler.ts
@@ -81,7 +81,7 @@ export class Reconciler<
 
   private announceQueue = new FIFO<{
     range: Range3d<SubspaceId>;
-    count: number;
+    isEmpty: boolean;
     wantResponse: boolean;
     covers: bigint | typeof COVERS_NONE;
   }>();
@@ -166,14 +166,14 @@ export class Reconciler<
     if (this.fingerprintScheme.isEqual(fingerprint, fingerprintOursFinal)) {
       this.announceQueue.push({
         range,
-        count: 0,
+        isEmpty: true,
         wantResponse: false,
         covers: BigInt(yourRangeCounter),
       });
     } else if (size <= SEND_ENTRIES_THRESHOLD) {
       this.announceQueue.push({
         range,
-        count: size,
+        isEmpty: size === 0,
         wantResponse: true,
         covers: BigInt(yourRangeCounter),
       });

diff --git a/src/wgps/wgps_messenger.ts b/src/wgps/wgps_messenger.ts
@@ -347,7 +347,7 @@ export class WgpsMessenger<
   private currentlyReceivingEntries: {
     namespace: NamespaceId;
     range: Range3d<SubspaceId>;
-    remaining: bigint;
+    expecting: boolean;
   } | undefined;
 
   private reconciliationPayloadIngester: PayloadIngester<
@@ -1020,7 +1020,7 @@ export class WgpsMessenger<
       // Then announce the entries.
       this.encoder.encode({
         kind: MsgKind.ReconciliationAnnounceEntries,
-        count: BigInt(pack.announcement.count),
+        isEmpty: pack.announcement.isEmpty,
         range: pack.announcement.range,
         wantResponse: pack.announcement.wantResponse,
         willSort: true,
@@ -1030,6 +1030,8 @@ export class WgpsMessenger<
       });
 
       // Then send the entries.
+      const finalEntry = pack.entries.at(-1);
+
       for (const entry of pack.entries) {
         this.encoder.encode({
           kind: MsgKind.ReconciliationSendEntry,
@@ -1042,6 +1044,7 @@ export class WgpsMessenger<
 
         this.encoder.encode({
           kind: MsgKind.ReconciliationTerminatePayload,
+          isFinal: entry === finalEntry,
         });
       }
     });
@@ -1333,7 +1336,7 @@ export class WgpsMessenger<
       case MsgKind.ReconciliationAnnounceEntries: {
         if (
           this.currentlyReceivingEntries &&
-          this.currentlyReceivingEntries.remaining > 0n
+          this.currentlyReceivingEntries.expecting
         ) {
           throw new WgpsMessageValidationError(
             "Never received the entries we were promised...",
@@ -1346,7 +1349,7 @@ export class WgpsMessenger<
 
         // Set the currently receiving namespace and range and expected count.
         this.currentlyReceivingEntries = {
-          remaining: message.count,
+          expecting: !message.isEmpty,
           namespace: reconciler.store.namespace,
           range: message.range,
         };
@@ -1387,7 +1390,7 @@ export class WgpsMessenger<
           );
         }
 
-        if (this.currentlyReceivingEntries.remaining <= 0n) {
+        if (!this.currentlyReceivingEntries.expecting) {
           throw new WgpsMessageValidationError(
             "Received entry for an announcement when we are not expecting any more!",
           );
@@ -1418,8 +1421,6 @@ export class WgpsMessenger<
           );
         }
 
-        this.currentlyReceivingEntries.remaining -= 1n;
-
         this.reconciliationPayloadIngester.target(
           message.entry.entry,
           message.entry.available === message.entry.entry.payloadLength,
@@ -1434,6 +1435,10 @@ export class WgpsMessenger<
         break;
       }
       case MsgKind.ReconciliationTerminatePayload: {
+        if (message.isFinal && this.currentlyReceivingEntries) {
+            this.currentlyReceivingEntries.expecting = false;
+        }
+
         const entryToRequestPayloadFor = this.reconciliationPayloadIngester
           .terminate();