@effect/cluster 0.52.2 → 0.52.4

package/src/Sharding.ts

@@ -11,7 +11,9 @@ import type { DurationInput } from "effect/Duration"
 import * as Effect from "effect/Effect"
 import * as Either from "effect/Either"
 import * as Equal from "effect/Equal"
+import type * as Exit from "effect/Exit"
 import * as Fiber from "effect/Fiber"
+import * as FiberHandle from "effect/FiberHandle"
 import * as FiberMap from "effect/FiberMap"
 import * as FiberRef from "effect/FiberRef"
 import * as FiberSet from "effect/FiberSet"
@@ -280,8 +282,19 @@ const make = Effect.gen(function*() {
           continue
         }
 
-        const acquired = yield* runnerStorage.acquire(selfAddress, unacquiredShards)
-        yield* Effect.ignore(storage.resetShards(acquired))
+        const oacquired = yield* runnerStorage.acquire(selfAddress, unacquiredShards).pipe(
+          Effect.timeoutOption(config.shardLockRefreshInterval)
+        )
+        if (Option.isNone(oacquired)) {
+          activeShardsLatch.unsafeOpen()
+          continue
+        }
+
+        const acquired = oacquired.value
+        yield* storage.resetShards(acquired).pipe(
+          Effect.ignore,
+          Effect.timeoutOption(config.shardLockRefreshInterval)
+        )
         for (const shardId of acquired) {
           if (MutableHashSet.has(releasingShards, shardId) || !MutableHashSet.has(selfShards, shardId)) {
             continue
@@ -349,27 +362,28 @@ const make = Effect.gen(function*() {
       Effect.forkIn(shardingScope)
     )
 
-    const releaseShardsLock = Effect.unsafeMakeSemaphore(1).withPermits(1)
-    const releaseShards = releaseShardsLock(
-      Effect.suspend(() =>
-        Effect.forEach(
-          releasingShards,
-          (shardId) =>
-            Effect.forEach(
-              entityManagers.values(),
-              (state) => state.manager.interruptShard(shardId),
-              { concurrency: "unbounded", discard: true }
-            ).pipe(
-              Effect.andThen(runnerStorage.release(selfAddress, shardId)),
-              Effect.annotateLogs({ runner: selfAddress }),
-              Effect.flatMap(() => {
-                MutableHashSet.remove(releasingShards, shardId)
-                return storage.unregisterShardReplyHandlers(shardId)
-              })
-            ),
-          { concurrency: "unbounded", discard: true }
-        )
+    const releaseShardsHandle = yield* FiberHandle.make()
+    const releaseShards = Effect.suspend(() =>
+      Effect.forEach(
+        releasingShards,
+        (shardId) =>
+          Effect.forEach(
+            entityManagers.values(),
+            (state) => state.manager.interruptShard(shardId),
+            { concurrency: "unbounded", discard: true }
+          ).pipe(
+            Effect.andThen(runnerStorage.release(selfAddress, shardId)),
+            Effect.annotateLogs({ runner: selfAddress }),
+            Effect.flatMap(() => {
+              MutableHashSet.remove(releasingShards, shardId)
+              return storage.unregisterShardReplyHandlers(shardId)
+            })
+          ),
+        { concurrency: "unbounded", discard: true }
       )
+    ).pipe(
+      Effect.repeat({ until: () => MutableHashSet.size(releasingShards) === 0 }),
+      FiberHandle.run(releaseShardsHandle, { onlyIfMissing: true })
     )
 
     // open the shard latch every poll interval
@@ -1247,7 +1261,9 @@ const make = Effect.gen(function*() {
         Effect.fiberIdWith((id) => {
           state.status = "closing"
           internalInterruptors.add(id)
-          return Effect.void
+          // if preemptive shutdown is enabled, we start shutting down Sharding
+          // too
+          return config.preemptiveShutdown ? shutdown() : Effect.void
         })
       )
 
@@ -1320,23 +1336,27 @@ const make = Effect.gen(function*() {
 
   // --- Finalization ---
 
-  yield* Scope.addFinalizerExit(
-    shardingScope,
-    Effect.fnUntraced(function*(exit) {
-      yield* Effect.logDebug("Shutting down", exit._tag === "Success" ? {} : exit.cause).pipe(
+  const shutdown = Effect.fnUntraced(function*(exit?: Exit.Exit<unknown, unknown>) {
+    if (exit) {
+      yield* Effect.logDebug("Shutting down", exit._tag === "Failure" ? exit.cause : {}).pipe(
         Effect.annotateLogs({
           package: "@effect/cluster",
           module: "Sharding"
         })
       )
-      const fiberId = yield* Effect.fiberId
-      MutableRef.set(isShutdown, true)
-      internalInterruptors.add(fiberId)
-      if (selfRunner) {
-        yield* Effect.ignore(runnerStorage.unregister(selfRunner.address))
-      }
-    })
-  )
+    }
+
+    const fiberId = yield* Effect.fiberId
+    internalInterruptors.add(fiberId)
+    if (isShutdown.current) return
+
+    MutableRef.set(isShutdown, true)
+    if (selfRunner) {
+      yield* Effect.ignore(runnerStorage.unregister(selfRunner.address))
+    }
+  })
+
+  yield* Scope.addFinalizerExit(shardingScope, shutdown)
 
   const activeEntityCount = Effect.gen(function*() {
     let count = 0

package/src/ShardingConfig.ts

@@ -63,6 +63,12 @@ export class ShardingConfig extends Context.Tag("@effect/cluster/ShardingConfig"
    * Shard lock expiration duration.
    */
   readonly shardLockExpiration: DurationInput
+  /**
+   * Start shutting down as soon as an Entity has started shutting down.
+   *
+   * Defaults to `true`.
+   */
+  readonly preemptiveShutdown: boolean
   /**
    * The default capacity of the mailbox for entities.
    */
@@ -118,6 +124,7 @@ export const defaults: ShardingConfig["Type"] = {
   runnerShardWeight: 1,
   shardsPerGroup: 300,
   shardGroups: ["default"],
+  preemptiveShutdown: true,
   shardLockRefreshInterval: Duration.seconds(10),
   shardLockExpiration: Duration.seconds(35),
   entityMailboxCapacity: 4096,
@@ -179,6 +186,10 @@ export const config: Config.Config<ShardingConfig["Type"]> = Config.all({
     Config.withDefault(defaults.shardsPerGroup),
     Config.withDescription("The number of shards to allocate per shard group.")
   ),
+  preemptiveShutdown: Config.boolean("preemptiveShutdown").pipe(
+    Config.withDefault(defaults.preemptiveShutdown),
+    Config.withDescription("Start shutting down as soon as an Entity has started shutting down.")
+  ),
   shardLockRefreshInterval: Config.duration("shardLockRefreshInterval").pipe(
     Config.withDefault(defaults.shardLockRefreshInterval),
     Config.withDescription("Shard lock refresh interval.")

package/src/SqlRunnerStorage.ts

@@ -242,10 +242,8 @@ export const make = Effect.fnUntraced(function*(options: {
         if (toAcquire.size === 0) {
           return acquiredShardIds
         }
-        const results = (yield* conn.executeUnprepared(`SELECT ${pgLocks(toAcquire)}`, [], undefined))[0] as Record<
-          string,
-          boolean
-        >
+        const rows = yield* conn.executeUnprepared(`SELECT ${pgLocks(toAcquire)}`, [], undefined)
+        const results = rows[0] as Record<string, boolean>
         for (const shardId in results) {
           if (results[shardId]) {
             acquiredShardIds.push(shardId)
@@ -449,7 +447,9 @@ export const make = Effect.fnUntraced(function*(options: {
           function*(_address, shardId) {
             const lockNum = lockNumbers.get(shardId)!
             const conn = yield* Resource.get(lockConnRef!)
-            const release = conn.executeRaw(`SELECT pg_advisory_unlock(${lockNum})`, [])
+            const release = conn.executeRaw(`SELECT pg_advisory_unlock(${lockNum})`, []).pipe(
+              Effect.timeoutOption(Duration.seconds(5))
+            )
             const check = conn.executeValues(
               `SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND granted = true AND pid = pg_backend_pid() AND objid = ${lockNum}`,
               []
@@ -460,6 +460,7 @@ export const make = Effect.fnUntraced(function*(options: {
               if (takenLocks.length === 0) return
             }
           },
+          Effect.timeout(config.shardLockExpiration),
           Effect.onError(() => Resource.refresh(lockConnRef!)),
           Effect.asVoid,
           PersistenceError.refail,
@@ -481,6 +482,7 @@ export const make = Effect.fnUntraced(function*(options: {
               if (takenLocks.length === 0 || takenLocks[0][0] !== 1) return
             }
           },
+          Effect.timeout(config.shardLockExpiration),
           Effect.onError(() => Resource.refresh(lockConnRef!)),
           Effect.asVoid,
           PersistenceError.refail,