effect 4.0.0-beta.99 → 4.0.0-rc.109

package/src/unstable/cluster/Sharding.ts

@@ -23,6 +23,7 @@ import * as Fiber from "../../Fiber.ts"
 import * as FiberMap from "../../FiberMap.ts"
 import { constant, flow } from "../../Function.ts"
 import * as HashRing from "../../HashRing.ts"
+import * as InternalRecord from "../../internal/record.ts"
 import * as Latch from "../../Latch.ts"
 import * as Layer from "../../Layer.ts"
 import * as MutableHashMap from "../../MutableHashMap.ts"
@@ -55,6 +56,7 @@ import { EntityReaper } from "./internal/entityReaper.ts"
 import { hashString } from "./internal/hash.ts"
 import { internalInterruptors } from "./internal/interruptors.ts"
 import { ResourceMap } from "./internal/resourceMap.ts"
+import { effectiveInterval } from "./internal/shardLock.ts"
 import * as Message from "./Message.ts"
 import * as MessageStorage from "./MessageStorage.ts"
 import * as Reply from "./Reply.ts"
@@ -122,7 +124,7 @@ export class Sharding extends Context.Service<Sharding, {
       entityId: string
     ) => RpcClient.RpcClient.From<
       Rpcs,
-      MailboxFull | AlreadyProcessingMessage | PersistenceError
+      MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
     >
   >
 
@@ -176,7 +178,7 @@ export class Sharding extends Context.Service<Sharding, {
     discard: boolean
   ) => Effect.Effect<
     void,
-    MailboxFull | AlreadyProcessingMessage | PersistenceError
+    MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
   >
 
   /**
@@ -217,6 +219,7 @@ interface EntityManagerState {
 
 const make = Effect.gen(function*() {
   const config = yield* ShardingConfig
+  const shardLockInterval = effectiveInterval(config)
   const shardGroups = shardGroupConfig(config)
   const getRunnerAddress = () => Option.getOrUndefined(config.runnerAddress)
   const clock = yield* Clock
@@ -237,9 +240,13 @@ const make = Effect.gen(function*() {
   const runnerStorage = yield* RunnerStorage
 
   const entityManagers = new Map<string, EntityManagerState>()
+  let entityRegistrationStartMillis: number | undefined
+  let entityRegistrationFallbackStartMillis: number | undefined
 
   const shardAssignments = MutableHashMap.empty<ShardId, RunnerAddress>()
   const selfShards = MutableHashSet.empty<ShardId>()
+  // open while shard lock storage is healthy
+  const shardLocksHealthyLatch = Latch.makeUnsafe(true)
 
   // the active shards are the ones that we have acquired the lock for
   const acquiredShards = MutableHashSet.empty<ShardId>()
@@ -277,6 +284,9 @@ const make = Effect.gen(function*() {
   // allow them to move to another runner.
 
   const releasingShards = MutableHashSet.empty<ShardId>()
+  // Shards whose entities must be force interrupted and locks fully released
+  // before normal reacquisition.
+  const forceReleasingShards = MutableHashSet.empty<ShardId>()
   const initialRunnerAddress = getRunnerAddress()
   if (initialRunnerAddress) {
     const selfAddress = initialRunnerAddress
@@ -286,37 +296,88 @@ const make = Effect.gen(function*() {
     })
 
     const releaseShardsMap = yield* FiberMap.make<ShardId>()
-    const releaseShard = Effect.fnUntraced(
-      function*(shardId: ShardId) {
-        const fibers = Arr.empty<Fiber.Fiber<void>>()
+    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*<E>(
+      shardIds: ReadonlyArray<ShardId>,
+      force: boolean,
+      release: Effect.Effect<void, E>
+    ) {
+      const fibers = Arr.empty<Fiber.Fiber<void>>()
+      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.forkScoped(state.manager.interruptShard(shardId, { force })))
         }
-        yield* Fiber.joinAll(fibers)
-        yield* runnerStorage.release(selfAddress, shardId)
+      }
+      yield* Fiber.joinAll(fibers)
+      yield* shardLocksHealthyLatch.await
+      yield* release
+      for (const shardId of shardIds) {
         MutableHashSet.remove(releasingShards, shardId)
+        MutableHashSet.remove(forceReleasingShards, shardId)
         yield* storage.unregisterShardReplyHandlers(shardId)
-      },
-      Effect.sandbox,
-      (effect, shardId) =>
+      }
+    })
+    const retryShardRelease =
+      (annotations: { readonly fiber: string; readonly shardId?: ShardId }) =>
+      <A, E, R>(effect: Effect.Effect<A, E, R>) =>
         effect.pipe(
+          Effect.sandbox,
           Effect.tapError((cause) =>
-            Effect.logDebug(`Could not release shard, retrying`, cause).pipe(
+            Effect.logDebug(`Could not release shards, retrying`, cause).pipe(
               Effect.annotateLogs({
                 module: "effect/cluster/Sharding",
-                fiber: "releaseShard",
                 runner: selfAddress,
-                shardId
-              })
+                ...annotations
+              }),
+              // Effect.eventually retries immediately, so space failures to
+              // avoid hot-looping while storage is unavailable.
+              Effect.andThen(Effect.sleep(50))
             )
           ),
-          Effect.eventually,
+          Effect.eventually
+        )
+    const releaseShard = Effect.fnUntraced(
+      function*(shardId: ShardId) {
+        yield* runShardRelease(
+          [shardId],
+          MutableHashSet.has(forceReleasingShards, shardId),
+          runnerStorage.release(selfAddress, shardId)
+        )
+      },
+      (effect, shardId) =>
+        effect.pipe(
+          retryShardRelease({ fiber: "releaseShard", shardId }),
           FiberMap.run(releaseShardsMap, shardId, { onlyIfMissing: true })
         )
     )
+    // The forced release ends with `runnerStorage.releaseAll`, which drops
+    // every lock held by this runner. Shards must not be reacquired through the
+    // normal path while it is pending, otherwise the bulk release wipes a lock
+    // the runner already considers acquired.
+    const forcedShardReleasePending = () => forcedShardReleaseRunning || MutableHashSet.size(forceReleasingShards) > 0
+    const releaseForcedShards = Effect.suspend(() => {
+      if (forcedShardReleaseRunning || MutableHashSet.size(forceReleasingShards) === 0) {
+        return Effect.void
+      }
+      forcedShardReleaseRunning = true
+      const shardIds = [...forceReleasingShards]
+      return runShardRelease(shardIds, true, runnerStorage.releaseAll(selfAddress)).pipe(
+        retryShardRelease({ fiber: "releaseForcedShards" }),
+        Effect.ensuring(Effect.sync(() => {
+          forcedShardReleaseRunning = false
+          activeShardsLatch.openUnsafe()
+        })),
+        Effect.forkIn(shardingScope),
+        Effect.asVoid
+      )
+    })
     const releaseShards = Effect.gen(function*() {
+      yield* releaseForcedShards
       for (const shardId of releasingShards) {
+        if (MutableHashSet.has(forceReleasingShards, shardId)) continue
         if (FiberMap.hasUnsafe(releaseShardsMap, shardId)) continue
         yield* releaseShard(shardId)
       }
@@ -336,11 +397,21 @@ const make = Effect.gen(function*() {
           MutableHashSet.add(releasingShards, shardId)
         }
 
-        if (MutableHashSet.size(releasingShards) > 0) {
+        if (MutableHashSet.size(releasingShards) > 0 || MutableHashSet.size(forceReleasingShards) > 0) {
           yield* Effect.forkIn(syncSingletons, shardingScope)
           yield* releaseShards
         }
 
+        if (!shardLocksHealthyLatch.isOpen()) {
+          continue
+        }
+
+        // Wait for the pending bulk release before reacquiring, so it cannot
+        // drop a lock acquired here. `releaseForcedShards` reopens the latch.
+        if (forcedShardReleasePending()) {
+          continue
+        }
+
         // if a shard has been assigned to this runner, we acquire it
         const unacquiredShards = MutableHashSet.empty<ShardId>()
         for (const shardId of selfShards) {
@@ -353,7 +424,7 @@ const make = Effect.gen(function*() {
         }
 
         const oacquired = yield* runnerStorage.acquire(selfAddress, unacquiredShards).pipe(
-          Effect.timeoutOption(config.shardLockRefreshInterval)
+          Effect.timeoutOption(shardLockInterval)
         )
         if (Option.isNone(oacquired)) {
           activeShardsLatch.openUnsafe()
@@ -363,10 +434,19 @@ const make = Effect.gen(function*() {
         const acquired = oacquired.value
         yield* storage.resetShards(acquired).pipe(
           Effect.ignore,
-          Effect.timeoutOption(config.shardLockRefreshInterval)
+          Effect.timeoutOption(shardLockInterval)
         )
+        // A forced release can start while `acquire` is in flight, so re-check
+        // it here as well as before acquiring.
+        const forcedReleasePending = forcedShardReleasePending()
         for (const shardId of acquired) {
-          if (MutableHashSet.has(releasingShards, shardId) || !MutableHashSet.has(selfShards, shardId)) {
+          if (
+            !shardLocksHealthyLatch.isOpen() ||
+            forcedReleasePending ||
+            MutableHashSet.has(releasingShards, shardId) ||
+            !MutableHashSet.has(selfShards, shardId)
+          ) {
+            MutableHashSet.add(releasingShards, shardId)
             continue
           }
           MutableHashSet.add(acquiredShards, shardId)
@@ -392,8 +472,52 @@ const make = Effect.gen(function*() {
       Effect.forkIn(shardingScope)
     )
 
-    // refresh the shard locks every `shardLockRefreshInterval`
-    yield* Effect.suspend(() =>
+    const markShardLocksUnhealthy = (cause: Cause.Cause<unknown>) =>
+      Effect.suspend(() => {
+        if (!shardLocksHealthyLatch.closeUnsafe()) return Effect.void
+
+        const affectedShards = MutableHashSet.fromIterable([...acquiredShards, ...releasingShards])
+        MutableHashSet.clear(selfShards)
+        MutableHashSet.clear(acquiredShards)
+        for (const shardId of affectedShards) {
+          MutableHashSet.add(releasingShards, shardId)
+          MutableHashSet.add(forceReleasingShards, shardId)
+        }
+        ClusterMetrics.shards.updateUnsafe(BigInt(0), Context.empty())
+        activeShardsLatch.openUnsafe()
+
+        return Effect.gen(function*() {
+          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 }
+              )
+            }
+          }
+          activeShardsLatch.openUnsafe()
+        })
+      })
+
+    const markShardLocksHealthy = Effect.suspend(() => {
+      if (!shardLocksHealthyLatch.openUnsafe()) return Effect.void
+
+      MutableHashSet.clear(selfShards)
+      MutableHashMap.forEach(shardAssignments, (runner, shardId) => {
+        if (isLocalRunner(runner)) {
+          MutableHashSet.add(selfShards, shardId)
+        }
+      })
+      activeShardsLatch.openUnsafe()
+      return Effect.logInfo("Shard lock storage has recovered")
+    })
+
+    const refreshShardLocks = Effect.suspend(() =>
       runnerStorage.refresh(selfAddress, [
         ...acquiredShards,
         ...releasingShards
@@ -421,12 +545,20 @@ const make = Effect.gen(function*() {
         times: 5,
         schedule: Schedule.spaced(50)
       }),
-      Effect.catchCause((cause) =>
-        Effect.logError("Could not refresh shard locks", cause).pipe(
-          Effect.andThen(clearSelfShards)
-        )
-      ),
-      Effect.repeat(Schedule.fixed(config.shardLockRefreshInterval)),
+      Effect.timeout(shardLockInterval),
+      Effect.catchCause(markShardLocksUnhealthy)
+    )
+
+    const probeShardLocks = runnerStorage.refresh(selfAddress, []).pipe(
+      Effect.timeout(shardLockInterval),
+      Effect.andThen(markShardLocksHealthy),
+      Effect.catchCause(() => Effect.void)
+    )
+
+    // Refresh shard locks at the lease-safe interval, or probe storage while
+    // lock ownership is uncertain.
+    yield* Effect.suspend(() => shardLocksHealthyLatch.isOpen() ? refreshShardLocks : probeShardLocks).pipe(
+      Effect.repeat(Schedule.fixed(shardLockInterval)),
       Effect.forever,
       Effect.forkIn(shardingScope)
     )
@@ -439,11 +571,6 @@ const make = Effect.gen(function*() {
     )
   }
 
-  const clearSelfShards = Effect.sync(() => {
-    MutableHashSet.clear(selfShards)
-    activeShardsLatch.openUnsafe()
-  })
-
   // --- Storage inbox ---
   //
   // Responsible for reading unprocessed messages from storage and sending them
@@ -461,12 +588,38 @@ const make = Effect.gen(function*() {
   const storageReadLock = Semaphore.makeUnsafe(1)
   const withStorageReadLock = storageReadLock.withPermits(1)
 
+  // --- Entity residency ---
+  //
+  // A runner-wide counter that bounds how many entities can be resident at
+  // the same time. Entity managers reserve a slot before spawning an entity,
+  // so the sequential storage read loop and concurrent volatile sends share
+  // the same cap.
+
+  const maxResidentEntities = config.maxResidentEntities
+  let residentEntityCount = 0
+  const residencyAtCapacityUnsafe = () =>
+    maxResidentEntities !== "unbounded" && residentEntityCount >= maxResidentEntities
+  const residency: EntityManager.Residency = {
+    admitUnsafe() {
+      if (residencyAtCapacityUnsafe()) return false
+      residentEntityCount++
+      return true
+    },
+    releaseUnsafe() {
+      residentEntityCount--
+      // a slot has freed up, so the storage read loop can admit messages for
+      // entities it previously had to skip
+      if (maxResidentEntities !== "unbounded") {
+        storageReadLatch.openUnsafe()
+      }
+    }
+  }
+
   if (storageEnabled && initialRunnerAddress) {
     const selfAddress = initialRunnerAddress
     const entityRegistrationTimeoutMillis = Duration.toMillis(
       Duration.fromInputUnsafe(config.entityRegistrationTimeout)
     )
-    const storageStartMillis = clock.currentTimeMillisUnsafe()
 
     yield* Effect.gen(function*() {
       yield* Effect.logDebug("Starting")
@@ -474,8 +627,14 @@ const make = Effect.gen(function*() {
 
       let index = 0
       let messages: Array<Message.Incoming<any>> = []
+      let deliveredThisRead = false
       const removableNotifications = new Set<PendingNotification>()
       const resetAddresses = MutableHashSet.empty<EntityAddress>()
+      const cappedAddresses = MutableHashSet.empty<EntityAddress>()
+
+      const markDelivered = Effect.sync(() => {
+        deliveredThisRead = true
+      })
 
       const processMessages = Effect.whileLoop({
         while: () => index < messages.length,
@@ -483,6 +642,16 @@ const make = Effect.gen(function*() {
         body: () => send
       })
 
+      const readAndProcess = Effect.fnUntraced(function*(options?: {
+        readonly limit?: number | undefined
+        readonly addresses?: ReadonlyArray<EntityAddress> | undefined
+      }) {
+        messages = yield* storage.unprocessedMessages(acquiredShards, options)
+        index = 0
+        yield* processMessages
+        return messages.length
+      })
+
       const send = Effect.catchCause(
         Effect.suspend(() => {
           const message = messages[index]
@@ -492,8 +661,15 @@ const make = Effect.gen(function*() {
           }
           const state = entityManagers.get(address.entityType)
           if (!state) {
-            const sinceStart = clock.currentTimeMillisUnsafe() - storageStartMillis
-            if (sinceStart < entityRegistrationTimeoutMillis) {
+            const now = clock.currentTimeMillisUnsafe()
+            const registrationStarted = entityRegistrationStartMillis !== undefined
+            const timeoutStartMillis = entityRegistrationStartMillis ??
+              (entityRegistrationFallbackStartMillis ??= now)
+            // If registration never starts, allow two intervals from the first missing read before failing.
+            const timeoutMillis = registrationStarted
+              ? entityRegistrationTimeoutMillis
+              : entityRegistrationTimeoutMillis * 2
+            if (now - timeoutStartMillis < timeoutMillis) {
               // reset address in the case that the entity is slow to register
               MutableHashSet.add(resetAddresses, address)
               return Effect.void
@@ -509,7 +685,7 @@ const make = Effect.gen(function*() {
           if (message._tag === "IncomingEnvelope" && isProcessing) {
             // If the message might affect a currently processing request, we
             // send it to the entity manager to be processed.
-            return state.manager.send(message)
+            return Effect.tap(state.manager.send(message), markDelivered)
           } else if (isProcessing || state.status === "closing") {
             // If the request is already processing, we skip it.
             // Or if the entity is closing, we skip all incoming messages.
@@ -521,6 +697,14 @@ const make = Effect.gen(function*() {
             entry.resume(Effect.void)
           }
 
+          // The runner is at entity capacity, so messages for entities that
+          // are not already resident are skipped. Their claims are released
+          // after the read, so they stay eligible for once a slot frees up.
+          if (residencyAtCapacityUnsafe() && !state.manager.isResidentUnsafe(address)) {
+            MutableHashSet.add(cappedAddresses, address)
+            return Effect.void
+          }
+
           // If the entity was resuming in another fiber, we add the message
           // id to the unprocessed set.
           const resumptionState = MutableHashMap.get(entityResumptionState, address)
@@ -531,7 +715,7 @@ const make = Effect.gen(function*() {
             }
             return Effect.void
           }
-          return state.manager.send(message)
+          return Effect.tap(state.manager.send(message), markDelivered)
         }),
         (cause) => {
           const message = messages[index]
@@ -548,13 +732,25 @@ const make = Effect.gen(function*() {
             }))
           }
           if (error.success._tag === "MailboxFull") {
+            const address = message.envelope.address
+            const state = entityManagers.get(address.entityType)
+            // A resident entity has a full per-entity mailbox; its messages
+            // are resumed from storage once there is capacity again.
             // MailboxFull can only happen for requests, so this cast is safe
-            return resumeEntityFromStorage(message as Message.IncomingRequest<any>)
+            if (message._tag === "IncomingRequest" && state?.manager.isResidentUnsafe(address)) {
+              return resumeEntityFromStorage(message as Message.IncomingRequest<any>)
+            }
+            // Otherwise the runner ran out of entity slots while the message
+            // was in flight; leave it in storage.
+            MutableHashSet.add(cappedAddresses, address)
+            return Effect.void
           }
           return Effect.void
         }
       )
 
+      const batchSize = Math.max(1, config.unprocessedMessageBatchSize)
+
       while (true) {
         // wait for the next poll interval, or if we get notified of a change
         yield* storageReadLatch.await
@@ -573,15 +769,58 @@ const make = Effect.gen(function*() {
           pendingNotifications.forEach((entry) => removableNotifications.add(entry))
         }
 
-        messages = yield* storage.unprocessedMessages(acquiredShards)
-        index = 0
-        yield* processMessages
+        deliveredThisRead = false
+        let readCount = 0
+        let fullBatch = false
+        let exhaustive = false
+
+        // First deliver messages for entities that are already resident, so
+        // they keep making progress even when the runner is at entity
+        // capacity.
+        if (residencyAtCapacityUnsafe()) {
+          const residentAddresses: Array<EntityAddress> = []
+          for (const state of entityManagers.values()) {
+            if (state.status === "closed") continue
+            for (const address of state.manager.residentAddressesUnsafe()) {
+              residentAddresses.push(address)
+            }
+          }
+          if (residentAddresses.length > 0) {
+            readCount = yield* readAndProcess({ limit: batchSize, addresses: residentAddresses })
+            fullBatch = readCount >= batchSize
+          }
+        }
+
+        // Then walk the remaining messages, spawning new entities while the
+        // runner has entity slots left.
+        if (!residencyAtCapacityUnsafe() && readCount < batchSize) {
+          const limit = batchSize - readCount
+          const read = yield* readAndProcess({ limit })
+          if (read >= limit) {
+            fullBatch = true
+          } else if (MutableHashSet.size(cappedAddresses) === 0) {
+            exhaustive = true
+          }
+        }
 
         if (removableNotifications.size > 0) {
-          removableNotifications.forEach(({ message, resume }) => {
-            pendingNotifications.delete(message.envelope.requestId)
-            resume(Effect.fail(new EntityNotAssignedToRunner({ address: message.envelope.address })))
-          })
+          // On an exhaustive read, a message that was not seen is no longer
+          // processable by this runner. At the entity cap, the messages are
+          // safely persisted and delivered once a slot frees up, so persisted
+          // senders succeed instead of failing. Otherwise the read was
+          // truncated by the batch size and the notifications stay registered
+          // for the immediately following read.
+          const capped = MutableHashSet.size(cappedAddresses) > 0 || residencyAtCapacityUnsafe()
+          if (exhaustive || capped) {
+            removableNotifications.forEach(({ message, resume }) => {
+              pendingNotifications.delete(message.envelope.requestId)
+              resume(
+                exhaustive
+                  ? Effect.fail(new EntityNotAssignedToRunner({ address: message.envelope.address }))
+                  : Effect.void
+              )
+            })
+          }
           removableNotifications.clear()
         }
         if (MutableHashSet.size(resetAddresses) > 0) {
@@ -593,9 +832,29 @@ const make = Effect.gen(function*() {
           }
           MutableHashSet.clear(resetAddresses)
         }
+        // Capture claims skipped because of the entity cap. They are reset in
+        // one storage operation after releasing the read lock.
+        let cappedAddressesToReset: Array<EntityAddress> | undefined
+        if (MutableHashSet.size(cappedAddresses) > 0) {
+          cappedAddressesToReset = Arr.fromIterable(cappedAddresses)
+          MutableHashSet.clear(cappedAddresses)
+        }
 
         // let the resuming entities check if they are done
         yield* storageReadLock.release(1)
+
+        if (cappedAddressesToReset !== undefined) {
+          yield* Effect.ignore(storage.resetAddresses(cappedAddressesToReset))
+        }
+
+        // A full batch means more messages could be waiting; start the next
+        // read immediately, as long as this read made progress (a full batch
+        // of skipped messages must not spin the loop). When the runner is at
+        // capacity, the next read waits for an entity to be removed or the
+        // next poll interval instead.
+        if (fullBatch && deliveredThisRead) {
+          storageReadLatch.openUnsafe()
+        }
       }
     }).pipe(
       Effect.scoped,
@@ -646,7 +905,7 @@ const make = Effect.gen(function*() {
         const resumptionState = Option.getOrThrow(MutableHashMap.get(entityResumptionState, address))
         let done = false
 
-        while (!done) {
+        while (!done) { // oxlint-disable-line no-unmodified-loop-condition
           // if the shard is no longer assigned to this runner, we stop
           if (!MutableHashSet.has(acquiredShards, address.shardId)) {
             return
@@ -834,18 +1093,40 @@ const make = Effect.gen(function*() {
     retries?: number
   ): Effect.Effect<
     void,
-    MailboxFull | AlreadyProcessingMessage | PersistenceError
+    MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
   > {
+    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: EntityNotAssignedToRunner) => {
+      if (!isPersisted) {
+        return shouldFail
+          ? Effect.fail(error)
+          : Effect.logDebug("Abandoning outgoing message during shutdown", message.envelope.address)
+      }
+      const persist = message._tag === "OutgoingRequest"
+        ? storage.saveRequest(message)
+        : storage.saveEnvelope(message)
+      return Effect.catchTag(persist, "MalformedMessage", Effect.die).pipe(
+        Effect.andThen(
+          shouldFail
+            ? Effect.fail(error)
+            : Effect.logWarning("Persisting outgoing message abandoned during shutdown", message.envelope.address)
+        )
+      )
+    }
     return Effect.catchFilter(
       Effect.suspend(() => {
         const address = message.envelope.address
-        const isPersisted = Context.get(
-          message._tag === "OutgoingRequest" ? message.annotations : message.rpc.annotations,
-          Persisted
-        )
         if (isPersisted && !storageEnabled) {
           return Effect.die("Sharding.sendOutgoing: Persisted messages require MessageStorage")
         }
+        if (shouldFail && MutableRef.get(isShutdown)) {
+          return Effect.fail(new EntityNotAssignedToRunner({ address }))
+        }
         const maybeRunner = MutableHashMap.get(shardAssignments, address.shardId)
         const runnerIsLocal = Option.isSome(maybeRunner) && isLocalRunner(maybeRunner.value)
         if (isPersisted) {
@@ -857,6 +1138,8 @@ const make = Effect.gen(function*() {
         }
         return runnerIsLocal
           ? sendLocal(message)
+          : discard
+          ? runnersService.notify({ address: maybeRunner, message, discard })
           : runnersService.send({ address: maybeRunner.value, message })
       }),
       (error) =>
@@ -864,6 +1147,15 @@ const make = Effect.gen(function*() {
           ? Result.succeed(error)
           : Result.fail(error),
       (error) => {
+        // Abandon the message during teardown: retrying would loop forever once the runner is shutting down
+        if (error._tag === "EntityNotAssignedToRunner") {
+          const targetManager = entityManagers.get(message.envelope.address.entityType)
+          const cannotRecover = MutableRef.get(isShutdown) ||
+            (targetManager !== undefined && targetManager.status !== "alive")
+          if (cannotRecover) {
+            return abandon(error)
+          }
+        }
         if (retries === 0) {
           return Effect.die(error)
         }
@@ -977,7 +1269,7 @@ const make = Effect.gen(function*() {
             if (newAssignments) {
               const runner = newAssignments[i]
               MutableHashMap.set(shardAssignments, shard, runner)
-              if (isLocalRunner(runner)) {
+              if (shardLocksHealthyLatch.isOpen() && isLocalRunner(runner)) {
                 MutableHashSet.add(selfShards, shard)
               }
             } else {
@@ -1031,7 +1323,7 @@ const make = Effect.gen(function*() {
     Entity<any, any>,
     (entityId: string) => RpcClient.RpcClient<
       any,
-      MailboxFull | AlreadyProcessingMessage
+      MailboxFull | AlreadyProcessingMessage | EntityNotAssignedToRunner
     >,
     never
   > = yield* ResourceMap.make(
@@ -1044,7 +1336,7 @@ const make = Effect.gen(function*() {
         flatten: true,
         onFromClient(options): Effect.Effect<
           void,
-          MailboxFull | AlreadyProcessingMessage | PersistenceError
+          MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
         > {
           const address = Context.getUnsafe(options.context, ClientAddressTag)
           switch (options.message._tag) {
@@ -1077,7 +1369,7 @@ const make = Effect.gen(function*() {
               if (!options.discard) {
                 const entry: ClientRequestEntry = {
                   rpc: rpc as any,
-                  context: fiber.currentContext,
+                  context: fiber.context,
                   message
                 }
                 clientRequests.set(id, entry)
@@ -1167,12 +1459,14 @@ const make = Effect.gen(function*() {
             return entity.protocol.requests.has(p as string)
           },
           get(target, p) {
-            if (p in target) {
+            if (Object.hasOwn(target, p)) {
               return target[p]
             } else if (!entity.protocol.requests.has(p as string)) {
               return undefined
             }
-            return target[p] = (payload: any, options?: {}) => clientFn(p as string, payload, options)
+            const method = (payload: any, options?: {}) => clientFn(p as string, payload, options)
+            InternalRecord.assignProperty(target, p, method)
+            return method
           }
         })
       }
@@ -1183,7 +1477,7 @@ const make = Effect.gen(function*() {
   const makeClient = <Type extends string, Rpcs extends Rpc.Any>(entity: Entity<Type, Rpcs>): Effect.Effect<
     (
       entityId: string
-    ) => RpcClient.RpcClient.From<Rpcs, MailboxFull | AlreadyProcessingMessage>
+    ) => RpcClient.RpcClient.From<Rpcs, MailboxFull | AlreadyProcessingMessage | EntityNotAssignedToRunner>
   > => clients.get(entity) as any
 
   const clientRespondDiscard = (_reply: Reply.Reply<any>) => Effect.void
@@ -1307,14 +1601,14 @@ const make = Effect.gen(function*() {
         ...options,
         storage,
         runnerAddress,
+        residency,
         sharding
       }).pipe(
-        Effect.provideContext(Context.mutate(services, (services) =>
-          services.pipe(
-            Context.add(EntityReaper, reaper),
-            Context.add(Scope.Scope, scope),
-            Context.add(Snowflake.Generator, snowflakeGen)
-          )))
+        Effect.provideContext(services.pipe(
+          Context.add(EntityReaper, reaper),
+          Context.add(Scope.Scope, scope),
+          Context.add(Snowflake.Generator, snowflakeGen)
+        ))
       ) as Effect.Effect<EntityManager.EntityManager>
       const state: EntityManagerState = {
         entity,
@@ -1335,6 +1629,7 @@ const make = Effect.gen(function*() {
       // register entities while storage is idle
       // this ensures message order is preserved
       yield* withStorageReadLock(Effect.sync(() => {
+        entityRegistrationStartMillis ??= clock.currentTimeMillisUnsafe()
         entityManagers.set(entity.type, state)
         if (entityManagerLatches.has(entity.type)) {
           entityManagerLatches.get(entity.type)!.openUnsafe()