effect 4.0.2 → 4.0.3

package/dist/cluster/Sharding.js

@@ -118,6 +118,13 @@ const make = /*#__PURE__*/Effect.gen(function* () {
   // the active shards are the ones that we have acquired the lock for
   const acquiredShards = MutableHashSet.empty();
   const activeShardsLatch = yield* Latch.make(false);
+  // open once shutdown starts handing every shard off
+  const handoffLatch = Latch.makeUnsafe(false);
+  // open once the handoff has released every shard lock
+  const handedOffLatch = Latch.makeUnsafe(false);
+  // serializes runner registration with unregistration, so a registration
+  // in flight cannot re-advertise a runner that is shutting down
+  const withRegistrationLock = Semaphore.makeUnsafe(1).withPermits(1);
   const events = yield* PubSub.unbounded();
   const getRegistrationEvents = Stream.fromPubSub(events);
   const isLocalRunner = address => {
@@ -131,6 +138,23 @@ const make = /*#__PURE__*/Effect.gen(function* () {
   function isEntityOnLocalShards(address) {
     return MutableHashSet.has(acquiredShards, address.shardId);
   }
+  function isAwaitingShardLock(shardId) {
+    return !MutableRef.get(isShutdown) && MutableHashSet.has(selfShards, shardId) && !MutableHashSet.has(acquiredShards, shardId);
+  }
+  // Requests for a shard assigned to this runner wait for its lock, for at most
+  // `shardLockExpiration`, instead of bouncing while the previous owner drains.
+  // Each notification opens the current latch and replaces it, so a waiter that
+  // checked its shard before a notification cannot miss it.
+  let shardLocksChanged = Latch.makeUnsafe(false);
+  const notifyShardLockWaiters = Effect.suspend(() => {
+    const latch = shardLocksChanged;
+    shardLocksChanged = Latch.makeUnsafe(false);
+    return latch.open;
+  });
+  const awaitShardLock = shardId => {
+    const loop = Effect.suspend(() => isAwaitingShardLock(shardId) ? Effect.andThen(shardLocksChanged.await, loop) : Effect.void);
+    return Effect.suspend(() => isAwaitingShardLock(shardId) ? Effect.asVoid(Effect.timeoutOption(loop, config.shardLockExpiration)) : Effect.void);
+  };
   yield* Scope.addFinalizer(shardingScope, Effect.logDebug("Shutdown complete").pipe(Effect.annotateLogs({
     module: "effect/cluster/Sharding"
   })));
@@ -141,6 +165,8 @@ const make = /*#__PURE__*/Effect.gen(function* () {
   // This should be shutdown last, when all entities have been shutdown, to
   // allow them to move to another runner.
   const releasingShards = MutableHashSet.empty();
+  // Hands every held shard off on shutdown, set when this is a runner
+  let handoff = Effect.void;
   // Shards whose entities must be force interrupted and locks fully released
   // before normal reacquisition.
   const forceReleasingShards = MutableHashSet.empty();
@@ -151,21 +177,40 @@ const make = /*#__PURE__*/Effect.gen(function* () {
       // the locks expire over time, so if this fails we ignore it
       return Effect.ignore(runnerStorage.releaseAll(selfAddress));
     });
-    const releaseShardsMap = yield* FiberMap.make();
-    let forcedShardReleaseRunning = false;
-    // Interrupt the shards' entities, wait for lock health, run the storage
-    // release, then clear the shards' bookkeeping.
-    const runShardRelease = Effect.fnUntraced(function* (shardIds, force, release) {
+    // Shard releases fork into a child scope that closes after the shutdown
+    // handoff, as forks into the closing sharding scope never run.
+    const releaseScope = yield* Scope.fork(shardingScope);
+    // Interrupt the shards' entities on every entity manager that is open
+    const forkInterruptShards = Effect.fnUntraced(function* (shardIds, force) {
       const fibers = Arr.empty();
       for (const shardId of shardIds) {
         for (const state of entityManagers.values()) {
           if (state.status === "closed") continue;
-          fibers.push(yield* Effect.forkScoped(state.manager.interruptShard(shardId, {
+          fibers.push(yield* Effect.forkIn(state.manager.interruptShard(shardId, {
             force
-          })));
+          }), releaseScope, {
+            startImmediately: true
+          }));
         }
       }
-      yield* Fiber.joinAll(fibers);
+      return fibers;
+    });
+    const releaseShardsMap = yield* FiberMap.make();
+    let forcedShardReleaseRunning = false;
+    // Called from the acquisition loop only, so no `acquire` is in flight
+    const checkHandedOff = () => {
+      if (handoffLatch.isOpen() && MutableHashSet.size(acquiredShards) === 0 && MutableHashSet.size(releasingShards) === 0 && !forcedShardReleasePending()) {
+        handedOffLatch.openUnsafe();
+      }
+    };
+    // Interrupt the shards' entities, wait for lock health, run the storage
+    // release, then clear the shards' bookkeeping.
+    const runShardRelease = Effect.fnUntraced(function* (shardIds, force, release) {
+      yield* Fiber.joinAll(yield* forkInterruptShards(shardIds, force));
+      // stop the singletons of shards that are no longer acquired, and wait
+      // for those on the released shards
+      yield* startSyncSingletons;
+      yield* awaitStoppingSingletons(shardIds);
       yield* shardLocksHealthyLatch.await;
       yield* release;
       for (const shardId of shardIds) {
@@ -179,6 +224,8 @@ const make = /*#__PURE__*/Effect.gen(function* () {
           interrupt: MutableRef.get(isShutdown)
         });
       }
+      // let the acquisition loop check whether the handoff has finished
+      if (handoffLatch.isOpen()) activeShardsLatch.openUnsafe();
     });
     const retryShardRelease = annotations => effect => effect.pipe(Effect.sandbox, Effect.tapError(cause => Effect.logDebug(`Could not release shards, retrying`, cause).pipe(Effect.annotateLogs({
       module: "effect/cluster/Sharding",
@@ -212,7 +259,7 @@ const make = /*#__PURE__*/Effect.gen(function* () {
       }), Effect.ensuring(Effect.sync(() => {
         forcedShardReleaseRunning = false;
         activeShardsLatch.openUnsafe();
-      })), Effect.forkIn(shardingScope), Effect.asVoid);
+      })), Effect.forkIn(releaseScope), Effect.asVoid);
     });
     const releaseShards = Effect.gen(function* () {
       yield* releaseForcedShards;
@@ -227,9 +274,11 @@ const make = /*#__PURE__*/Effect.gen(function* () {
       while (true) {
         yield* activeShardsLatch.await;
         activeShardsLatch.closeUnsafe();
-        // if a shard is no longer assigned to this runner, we release it
+        // if a shard is no longer assigned to this runner, or the runner is
+        // handing every shard off, we release it
+        const handingOff = handoffLatch.isOpen();
         for (const shardId of acquiredShards) {
-          if (MutableHashSet.has(selfShards, shardId)) continue;
+          if (!handingOff && MutableHashSet.has(selfShards, shardId)) continue;
           MutableHashSet.remove(acquiredShards, shardId);
           MutableHashSet.add(releasingShards, shardId);
         }
@@ -237,6 +286,7 @@ const make = /*#__PURE__*/Effect.gen(function* () {
           yield* Effect.forkIn(syncSingletons, shardingScope);
           yield* releaseShards;
         }
+        checkHandedOff();
         if (!shardLocksHealthyLatch.isOpen()) {
           continue;
         }
@@ -245,6 +295,10 @@ const make = /*#__PURE__*/Effect.gen(function* () {
         if (forcedShardReleasePending()) {
           continue;
         }
+        // a runner that is shutting down takes no new shards
+        if (isShutdown.current) {
+          continue;
+        }
         // if a shard has been assigned to this runner, we acquire it
         const unacquiredShards = MutableHashSet.empty();
         for (const shardId of selfShards) {
@@ -272,12 +326,14 @@ const make = /*#__PURE__*/Effect.gen(function* () {
           MutableHashSet.add(acquiredShards, shardId);
         }
         if (acquired.length > 0) {
+          yield* notifyShardLockWaiters;
           yield* storageReadLatch.open;
           yield* Effect.forkIn(syncSingletons, shardingScope);
           // update metrics
           ClusterMetrics.shards.updateUnsafe(BigInt(MutableHashSet.size(acquiredShards)), Context.empty());
         }
-        yield* Effect.sleep(1000);
+        // the shutdown handoff releases the shards acquired here without waiting
+        yield* Effect.raceFirst(Effect.sleep(1000), handoffLatch.await);
         activeShardsLatch.openUnsafe();
       }
     }).pipe(Effect.catchCause(cause => Effect.logWarning("Could not acquire/release shards", cause)), Effect.repeat(Schedule.spaced(config.entityMessagePollInterval)), Effect.annotateLogs({
@@ -297,20 +353,12 @@ const make = /*#__PURE__*/Effect.gen(function* () {
       ClusterMetrics.shards.updateUnsafe(BigInt(0), Context.empty());
       activeShardsLatch.openUnsafe();
       return Effect.gen(function* () {
+        yield* notifyShardLockWaiters;
         yield* Effect.logError("Shard lock storage is unhealthy", cause);
         yield* Effect.forkIn(syncSingletons, shardingScope, {
           startImmediately: true
         });
-        for (const shardId of affectedShards) {
-          for (const state of entityManagers.values()) {
-            if (state.status === "closed") continue;
-            yield* Effect.forkIn(state.manager.interruptShard(shardId, {
-              force: true
-            }), shardingScope, {
-              startImmediately: true
-            });
-          }
-        }
+        yield* forkInterruptShards(affectedShards, true);
         activeShardsLatch.openUnsafe();
       });
     });
@@ -355,6 +403,14 @@ const make = /*#__PURE__*/Effect.gen(function* () {
     yield* Effect.suspend(() => shardLocksHealthyLatch.isOpen() ? refreshShardLocks : probeShardLocks).pipe(Effect.repeat(Schedule.fixed(shardLockInterval)), Effect.forever, Effect.forkIn(shardingScope));
     // open the shard latch every poll interval
     yield* activeShardsLatch.open.pipe(Effect.delay(config.entityMessagePollInterval), Effect.forever, Effect.forkIn(shardingScope));
+    // Release every held shard through the per-shard path. Their entities are
+    // forced to stop, so the handoff does not wait out graceful termination.
+    // This is best effort: the wait is bounded, and if it times out, the final
+    // `releaseAll` attempts to release the remaining locks.
+    handoff = Effect.suspend(() => {
+      if (handoffLatch.openUnsafe()) activeShardsLatch.openUnsafe();
+      return forkInterruptShards([...acquiredShards, ...releasingShards], true);
+    }).pipe(Effect.andThen(Effect.timeoutOption(handedOffLatch.await, shardLockInterval)), Effect.asVoid);
   }
   // --- Storage inbox ---
   //
@@ -717,28 +773,34 @@ const make = /*#__PURE__*/Effect.gen(function* () {
     }), (effect, address) => Effect.ensuring(effect, Effect.sync(() => MutableHashMap.remove(entityResumptionState, address))), Effect.forkIn(shardingScope));
   }
   // --- Sending messages ---
-  const sendLocal = message => Effect.suspend(function loop() {
+  const sendLocal = message => {
     const address = message.envelope.address;
-    if (!isEntityOnLocalShards(address)) {
-      return Effect.fail(new EntityNotAssignedToRunner({
-        address
-      }));
-    }
-    const state = entityManagers.get(address.entityType);
-    if (!state) {
-      return Effect.flatMap(waitForEntityManager(address.entityType), loop);
-    } else if (state.status === "closed" || state.status === "closing" && message._tag === "IncomingRequest") {
-      // if we are shutting down, we don't accept new requests
-      return Effect.fail(new EntityNotAssignedToRunner({
-        address
-      }));
-    }
-    return message._tag === "IncomingRequest" || message._tag === "IncomingEnvelope" ? state.manager.send(message) : runnersService.sendLocal({
-      message,
-      send: state.manager.sendLocal,
-      simulateRemoteSerialization: config.simulateRemoteSerialization
+    const isRequest = message._tag === "IncomingRequest" || message._tag === "OutgoingRequest";
+    const send = Effect.suspend(function loop() {
+      if (!isEntityOnLocalShards(address) && (
+      // Entities draining from a released shard still accept interrupts and acks
+      isRequest || !MutableHashSet.has(releasingShards, address.shardId))) {
+        return Effect.fail(new EntityNotAssignedToRunner({
+          address
+        }));
+      }
+      const state = entityManagers.get(address.entityType);
+      if (!state) {
+        return Effect.flatMap(waitForEntityManager(address.entityType), loop);
+      } else if (state.status === "closed" || state.status === "closing" && message._tag === "IncomingRequest") {
+        // if we are shutting down, we don't accept new requests
+        return Effect.fail(new EntityNotAssignedToRunner({
+          address
+        }));
+      }
+      return message._tag === "IncomingRequest" || message._tag === "IncomingEnvelope" ? state.manager.send(message) : runnersService.sendLocal({
+        message,
+        send: state.manager.sendLocal,
+        simulateRemoteSerialization: config.simulateRemoteSerialization
+      });
     });
-  });
+    return isRequest ? Effect.andThen(awaitShardLock(address.shardId), send) : send;
+  };
   const pendingNotifications = new Map();
   const notifyLocal = (message, discard, options) => Effect.suspend(function loop() {
     const address = message.envelope.address;
@@ -793,7 +855,10 @@ const make = /*#__PURE__*/Effect.gen(function* () {
       storageOnly: !isLocal
     });
   });
-  function sendOutgoing(message, discard, retries) {
+  function sendOutgoing(message, discard, retries,
+  // The runner serving the request, for its interrupts and acks. Later
+  // attempts fall back to the current shard assignment.
+  runner) {
     const isPersisted = Context.get(message._tag === "OutgoingRequest" ? message.annotations : message.rpc.annotations, Persisted);
     const shouldFail = !discard && (message._tag === "OutgoingRequest" || message.envelope._tag === "AckChunk");
     const abandon = error => {
@@ -818,7 +883,7 @@ const make = /*#__PURE__*/Effect.gen(function* () {
           address
         }));
       }
-      const maybeRunner = MutableHashMap.get(shardAssignments, address.shardId);
+      const maybeRunner = runner !== undefined ? Option.some(runner) : MutableHashMap.get(shardAssignments, address.shardId);
       const runnerIsLocal = Option.isSome(maybeRunner) && isLocalRunner(maybeRunner.value);
       if (isPersisted) {
         return runnerIsLocal ? notifyLocal(message, discard) : runnersService.notify({
@@ -830,6 +895,9 @@ const make = /*#__PURE__*/Effect.gen(function* () {
         return Effect.fail(new EntityNotAssignedToRunner({
           address
         }));
+      } else if (message._tag === "OutgoingRequest") {
+        const entry = clientRequests.get(message.envelope.requestId);
+        if (entry) entry.runner = maybeRunner.value;
       }
       return runnerIsLocal ? sendLocal(message) : discard ? runnersService.notify({
         address: maybeRunner,
@@ -886,9 +954,12 @@ const make = /*#__PURE__*/Effect.gen(function* () {
     while (true) {
       // Ensure the current runner is registered
       if (selfRunner && !isShutdown.current && !MutableHashMap.has(allRunners, selfRunner)) {
-        yield* Effect.logDebug("Registering runner", selfRunner);
-        const machineId = yield* withTimeout(runnerStorage.register(selfRunner, true));
-        yield* snowflakeGen.setMachineId(machineId);
+        yield* withRegistrationLock(Effect.gen(function* () {
+          if (isShutdown.current) return;
+          yield* Effect.logDebug("Registering runner", selfRunner);
+          const machineId = yield* withTimeout(runnerStorage.register(selfRunner, true));
+          yield* snowflakeGen.setMachineId(machineId);
+        }));
       }
       const runners = yield* withTimeout(runnerStorage.getRunners);
       let changed = false;
@@ -957,6 +1028,7 @@ const make = /*#__PURE__*/Effect.gen(function* () {
         });
         yield* Effect.logDebug("New shard assignments", selfShards);
         activeShardsLatch.openUnsafe();
+        yield* notifyShardLockWaiters;
         // update metrics
         if (selfRunner) {
           ClusterMetrics.runnersHealthy.updateUnsafe(BigInt(MutableHashSet.has(healthyRunners, selfRunner) ? 1 : 0), Context.empty());
@@ -1040,7 +1112,7 @@ const make = /*#__PURE__*/Effect.gen(function* () {
                   replyId: entry.lastChunkId
                 }),
                 rpc: entry.rpc
-              }), false);
+              }), false, undefined, entry.runner);
             }
           case "Interrupt":
             {
@@ -1066,7 +1138,7 @@ const make = /*#__PURE__*/Effect.gen(function* () {
                   requestId
                 }),
                 rpc: entry.rpc
-              }), false, 3));
+              }), false, 3, entry.runner));
             }
         }
         return Effect.void;
@@ -1139,6 +1211,41 @@ const make = /*#__PURE__*/Effect.gen(function* () {
   const singletons = new Map();
   const singletonFibers = yield* FiberMap.make();
   const withSingletonLock = Semaphore.makeUnsafe(1).withPermits(1);
+  // Singletons that are stopping. They count as running, so they are not
+  // restarted, and shard releases wait for them. Their teardown is awaited
+  // without the singleton lock, as it can register or stop other singletons.
+  const stoppingSingletons = MutableHashMap.empty();
+  const isSingletonRegistered = address => {
+    const map = singletons.get(address.shardId);
+    return map !== undefined && MutableHashMap.has(map, address);
+  };
+  // Called with the singleton lock held. Starts stopping the singleton and
+  // returns the wait for its teardown, to run without the lock.
+  const stopSingletonUnsafe = (address, activeTeardown) => {
+    const stopping = MutableHashMap.get(stoppingSingletons, address);
+    if (Option.isSome(stopping)) return Fiber.join(stopping.value);
+    const fiber = FiberMap.getUnsafe(singletonFibers, address);
+    if (Option.isNone(fiber)) return Effect.void;
+    const interrupt = Fiber.interrupt(fiber.value);
+    const stop = Effect.runForkWith(services)(activeTeardown ? ActiveTeardown.aroundShard(address.shardId, interrupt) : interrupt);
+    MutableHashMap.set(stoppingSingletons, address, stop);
+    stop.addObserver(() => {
+      MutableHashMap.remove(stoppingSingletons, address);
+      // a registration made while the previous singleton was stopping
+      if (isSingletonRegistered(address)) runFork(syncSingletons);
+    });
+    return Fiber.join(stop);
+  };
+  // Waits for the singletons of the given shards that are still stopping
+  const awaitStoppingSingletons = shardIds => Effect.suspend(() => {
+    const stops = Arr.empty();
+    MutableHashMap.forEach(stoppingSingletons, (stop, address) => {
+      if (shardIds.some(shardId => Equal.equals(shardId, address.shardId))) {
+        stops.push(stop);
+      }
+    });
+    return Fiber.joinAll(stops);
+  });
   const registerSingleton = Effect.fnUntraced(function* (name, run, options) {
     const shardGroup = options?.shardGroup ?? "default";
     const address = new SingletonAddress({
@@ -1159,25 +1266,29 @@ const make = /*#__PURE__*/Effect.gen(function* () {
     yield* PubSub.publish(events, SingletonRegistered({
       address
     }));
-    // start if we are on the right shard
-    if (MutableHashSet.has(acquiredShards, address.shardId)) {
+    // start if we are on the right shard, unless a previous registration of
+    // this singleton is still stopping
+    if (MutableHashSet.has(acquiredShards, address.shardId) && !MutableHashMap.has(stoppingSingletons, address)) {
       yield* Effect.logDebug("Starting singleton", address);
       yield* FiberMap.run(singletonFibers, address, wrappedRun);
     }
-    yield* Effect.addFinalizer(() => {
-      const map = singletons.get(address.shardId);
-      MutableHashMap.remove(map, address);
-      return FiberMap.remove(singletonFibers, address);
-    });
+    yield* Effect.addFinalizer(() => withSingletonLock(Effect.sync(() => {
+      MutableHashMap.remove(singletons.get(address.shardId), address);
+      return stopSingletonUnsafe(address, false);
+    })).pipe(Effect.flatten));
   }, withSingletonLock);
-  const syncSingletons = withSingletonLock(Effect.gen(function* () {
+  // Starts and stops singletons to match the acquired shards, returning the
+  // waits for the stopped singletons
+  const startSyncSingletons = withSingletonLock(Effect.gen(function* () {
+    const stops = Arr.empty();
     for (const [shardId, map] of singletons) {
       for (const [address, run] of map) {
-        const running = FiberMap.hasUnsafe(singletonFibers, address);
+        const stopping = MutableHashMap.has(stoppingSingletons, address);
+        const running = stopping || FiberMap.hasUnsafe(singletonFibers, address);
         const shouldBeRunning = MutableHashSet.has(acquiredShards, shardId);
         if (running && !shouldBeRunning) {
-          yield* Effect.logDebug("Stopping singleton", address);
-          yield* ActiveTeardown.aroundShard(address.shardId, FiberMap.remove(singletonFibers, address));
+          if (!stopping) yield* Effect.logDebug("Stopping singleton", address);
+          stops.push(stopSingletonUnsafe(address, true));
         } else if (!running && shouldBeRunning) {
           yield* Effect.logDebug("Starting singleton", address);
           yield* FiberMap.run(singletonFibers, address, run);
@@ -1185,6 +1296,12 @@ const make = /*#__PURE__*/Effect.gen(function* () {
       }
     }
     ClusterMetrics.singletons.updateUnsafe(BigInt(yield* FiberMap.size(singletonFibers)), Context.empty());
+    return stops;
+  }));
+  const syncSingletons = Effect.flatMap(startSyncSingletons,
+  // wait for the stopped singletons without the lock
+  stops => Effect.all(stops, {
+    discard: true
   }));
   // --- Entities ---
   const reaper = yield* EntityReaper;
@@ -1283,11 +1400,13 @@ const make = /*#__PURE__*/Effect.gen(function* () {
     }
     if (isShutdown.current) return;
     MutableRef.set(isShutdown, true);
+    yield* notifyShardLockWaiters;
     if (selfRunner) {
-      yield* Effect.ignore(runnerStorage.unregister(selfRunner.address));
+      yield* Effect.ignore(withRegistrationLock(runnerStorage.unregister(selfRunner.address)));
     }
   });
-  yield* Scope.addFinalizerExit(shardingScope, shutdown);
+  // hand the shards off before the loops stop
+  yield* Scope.addFinalizerExit(shardingScope, exit => Effect.andThen(shutdown(exit), handoff));
   const activeEntityCount = Effect.gen(function* () {
     let count = 0;
     for (const state of entityManagers.values()) {