effect 4.0.2 → 4.0.4

package/src/cluster/internal/entityManager.ts

@@ -86,6 +86,10 @@ export interface Residency {
   readonly releaseUnsafe: () => void
 }
 
+interface RequestTransaction {
+  settle?: ((outcome: Exit.Exit<unknown, unknown>) => Effect.Effect<void>) | undefined
+}
+
 // Represents the entities managed by this entity manager
 /**
  * @internal
@@ -100,10 +104,13 @@ export type EntityState = {
     /** Excludes requests awaiting their first dispatch from replay. */
     delivered: boolean
     sentExit: boolean
+    /** Treat early termination interrupts like shutdown interrupts. */
+    terminating?: boolean | undefined
     lastSentChunk: Option.Option<Reply.Chunk<Rpc.Any>>
     sequence: number
     /** Set when the request should not outlive its caller. */
     callerScope?: Scope.Scope | undefined
+    transaction?: RequestTransaction | undefined
   }>
   lastActiveCheck: number
   write: RpcServer.RpcServer<any>["write"]
@@ -152,9 +159,13 @@ export const make = Effect.fnUntraced(function*<
   entityRpcs.set(KeepAliveRpc._tag, KeepAliveRpc as any)
 
   const activeServers = new Map<EntityId, EntityState>()
+  // Entities finishing their in-flight requests after removal, which still
+  // accept interrupts and acks
+  const drainingServers = new Map<EntityId, EntityState>()
   const serverCloseLatches = new Map<EntityAddress, {
     readonly closed: Latch.Latch
     readonly force: Latch.Latch
+    closing: boolean
   }>()
   const processedRequestIds = new Set<Snowflake.Snowflake>()
 
@@ -181,7 +192,8 @@ export const make = Effect.fnUntraced(function*<
     const keepAliveLatch = Latch.makeUnsafe()
     const closeLatches = {
       closed: Latch.makeUnsafe(),
-      force: Latch.makeUnsafe()
+      force: Latch.makeUnsafe(),
+      closing: false
     }
 
     yield* Scope.addFinalizer(
@@ -264,22 +276,23 @@ export const make = Effect.fnUntraced(function*<
                 // They will be retried when the entity is restarted.
                 // Also, if the request is uninterruptible, we ignore the
                 // interrupt.
+                const persisted = storageEnabled && Context.get(request.message.annotations, Persisted)
                 if (
-                  storageEnabled &&
-                  Context.get(request.message.annotations, Persisted) &&
+                  persisted &&
                   Exit.hasInterrupts(response.exit) &&
-                  (isShuttingDown || isUninterruptibleForServer(request.message.annotations))
+                  (isShuttingDown || request.terminating || isUninterruptibleForServer(request.message.annotations))
                 ) {
-                  if (!isShuttingDown) {
+                  if (!isShuttingDown && !request.terminating) {
                     request.sentExit = false
                     return server.write(0, requestEnvelope(request), requestWriteOptions(request)).pipe(
+                      Effect.setContext(handlerContext),
                       Effect.forkIn(handlerScope)
                     )
                   }
                   activeRequests.delete(Snowflake.Snowflake(response.requestId))
                   return options.storage.unregisterReplyHandler(request.message.envelope.requestId)
                 }
-                return retryRespond(
+                const save = retryRespond(
                   4,
                   Effect.suspend(() =>
                     request.message.respond(
@@ -290,22 +303,53 @@ export const make = Effect.fnUntraced(function*<
                       })
                     )
                   )
-                ).pipe(
-                  Effect.flatMap(() => {
-                    if (storageEnabled) {
-                      processedRequestIds.add(request.message.envelope.requestId)
-                    }
-                    activeRequests.delete(Snowflake.Snowflake(response.requestId))
+                )
+                const complete = Effect.sync(() => {
+                  if (storageEnabled) {
+                    processedRequestIds.add(request.message.envelope.requestId)
+                  }
+                  activeRequests.delete(Snowflake.Snowflake(response.requestId))
 
-                    // Start the idle timer when the last request completes.
-                    if (activeRequests.size === 0) {
-                      state.lastActiveCheck = clock.currentTimeMillisUnsafe()
-                    }
+                  // Start the idle timer when the last request completes.
+                  if (activeRequests.size === 0) {
+                    state.lastActiveCheck = clock.currentTimeMillisUnsafe()
+                  }
+                })
+                const respond = Effect.orDie(Effect.andThen(save, complete))
+                const transaction = request.transaction
+                if (!transaction) return respond
+                const exit = response.exit
+
+                if (!persisted) {
+                  if (Exit.isSuccess(exit)) return respond
+                  transaction.settle = () => respond
+                  return Effect.void
+                }
 
-                    return Effect.void
-                  }),
-                  Effect.orDie
-                )
+                if (Exit.isSuccess(exit)) {
+                  transaction.settle = (outcome) =>
+                    Exit.isSuccess(outcome) ? complete : Effect.flatMap(
+                      hasStoredExit(request.message.envelope.requestId),
+                      (committed) => {
+                        if (!committed) return restartFrom(outcome.cause)
+                        // Commit succeeded despite the error; let callers read the saved reply.
+                        return Effect.andThen(
+                          complete,
+                          options.storage.unregisterReplyHandler(request.message.envelope.requestId)
+                        )
+                      }
+                    )
+                  return Effect.orDie(save)
+                }
+                transaction.settle = (outcome) =>
+                  // Transaction wrappers can add defects, but not typed errors; interrupts may be repeated.
+                  Exit.isSuccess(outcome) || !outcome.cause.reasons.some((reason) =>
+                      Cause.isDieReason(reason) &&
+                      !exit.cause.reasons.some((r) => Cause.isDieReason(r) && r.defect === reason.defect)
+                    )
+                    ? respond
+                    : restartFrom(outcome.cause)
+                return Effect.void
               }
               case "Chunk": {
                 const request = activeRequests.get(Snowflake.Snowflake(response.requestId))
@@ -330,10 +374,7 @@ export const make = Effect.fnUntraced(function*<
                 ))
               }
               case "Defect": {
-                if (!isActive()) return endLatch.open
-                const rebuild = writeRef.rebuildUnsafe({ from: server.write, prepare: replay })
-                if (!rebuild) return Effect.void
-                return Effect.forkIn(restart(Cause.die(response.defect), rebuild), managerScope)
+                return restartFrom(Cause.die(response.defect))
               }
               case "ClientEnd": {
                 return endLatch.open
@@ -345,6 +386,21 @@ export const make = Effect.fnUntraced(function*<
           Effect.setContext(Context.merge(handlerContext, handlers))
         )
 
+        // Replay if the stored reply cannot be read.
+        const hasStoredExit = (requestId: Snowflake.Snowflake): Effect.Effect<boolean> =>
+          options.storage.repliesForUnfiltered([requestId]).pipe(
+            Effect.map((replies) => replies.some((reply) => reply._tag === "WithExit")),
+            Effect.catchCause(() => Effect.succeed(false))
+          )
+
+        // Do not inherit the failed handler's transaction context.
+        const restartFrom = (cause: Cause.Cause<unknown>): Effect.Effect<void> => {
+          if (!isActive()) return endLatch.open
+          const rebuild = writeRef.rebuildUnsafe({ from: server.write, prepare: replay })
+          if (!rebuild) return Effect.void
+          return Effect.forkIn(Effect.setContext(restart(cause, rebuild), handlerContext), managerScope)
+        }
+
         yield* Scope.addFinalizer(
           handlerScope,
           Effect.sync(() => {
@@ -415,15 +471,32 @@ export const make = Effect.fnUntraced(function*<
       Effect.suspend(() => {
         activeServers.delete(address.entityId)
         retired.openUnsafe()
+        closeLatches.closing = true
         acquireEntity(address)
+        drainingServers.set(address.entityId, state)
         return Effect.raceFirst(
-          state.write(0, { _tag: "Eof" }).pipe(
+          Effect.gen(function*() {
+            for (const [requestId, request] of activeRequests) {
+              if (
+                !request.delivered || request.sentExit ||
+                !Context.get(request.message.annotations, ClusterSchema.InterruptOnTermination) ||
+                isUninterruptibleForServer(request.message.annotations)
+              ) continue
+              request.terminating = true
+              yield* state.write(0, { _tag: "Interrupt", requestId: requestId as any, interruptors: [] })
+            }
+            yield* state.write(0, { _tag: "Eof" })
+          }).pipe(
             Effect.andThen(endLatch.await),
             Effect.timeoutOption(config.entityTerminationTimeout),
             Effect.interruptible
           ),
           Effect.interruptible(closeLatches.force.await)
-        )
+        ).pipe(Effect.ensuring(Effect.sync(() => {
+          if (drainingServers.get(address.entityId) === state) {
+            drainingServers.delete(address.entityId)
+          }
+        })))
       })
     )
     if (!options.sharding.hasShardId(address.shardId)) {
@@ -458,9 +531,38 @@ export const make = Effect.fnUntraced(function*<
     Effect.forkIn(managerScope)
   )
 
+  function sendEnvelope(server: EntityState, message: Message.IncomingEnvelope): Effect.Effect<void> {
+    const entry = server.activeRequests.get(message.envelope.requestId)
+    if (!entry) {
+      return Effect.void
+    } else if (
+      message.envelope._tag === "AckChunk" &&
+      Option.isSome(entry.lastSentChunk) &&
+      message.envelope.replyId !== entry.lastSentChunk.value.id
+    ) {
+      return Effect.void
+    }
+    return server.write(
+      0,
+      message.envelope._tag === "AckChunk"
+        ? { _tag: "Ack", requestId: message.envelope.requestId as any }
+        : {
+          _tag: "Interrupt",
+          requestId: message.envelope.requestId as any,
+          interruptors: []
+        }
+    )
+  }
+
   function sendLocal<R extends Rpc.Any>(
     message: Message.IncomingLocal<R>
   ): Effect.Effect<void, EntityNotAssignedToRunner | MailboxFull | AlreadyProcessingMessage> {
+    if (message._tag === "IncomingEnvelope") {
+      const draining = drainingServers.get(message.envelope.address.entityId)
+      if (draining?.activeRequests.has(message.envelope.requestId)) {
+        return sendEnvelope(draining, message)
+      }
+    }
     return Effect.provideService(
       Effect.flatMap(
         entities.get(message.envelope.address),
@@ -557,26 +659,7 @@ export const make = Effect.fnUntraced(function*<
               return server.write(0, requestEnvelope(entry), requestWriteOptions(entry))
             }
             case "IncomingEnvelope": {
-              const entry = server.activeRequests.get(message.envelope.requestId)
-              if (!entry) {
-                return Effect.void
-              } else if (
-                message.envelope._tag === "AckChunk" &&
-                Option.isSome(entry.lastSentChunk) &&
-                message.envelope.replyId !== entry.lastSentChunk.value.id
-              ) {
-                return Effect.void
-              }
-              return server.write(
-                0,
-                message.envelope._tag === "AckChunk"
-                  ? { _tag: "Ack", requestId: message.envelope.requestId as any }
-                  : {
-                    _tag: "Interrupt",
-                    requestId: message.envelope.requestId as any,
-                    interruptors: []
-                  }
-              )
+              return sendEnvelope(server, message)
             }
           }
         }
@@ -596,11 +679,24 @@ export const make = Effect.fnUntraced(function*<
     })
 
   const requestWriteOptions = (
-    entry: { readonly message: Message.IncomingRequestLocal<any>; readonly callerScope?: Scope.Scope | undefined }
+    entry: {
+      readonly message: Message.IncomingRequestLocal<any>
+      readonly callerScope?: Scope.Scope | undefined
+      transaction?: RequestTransaction | undefined
+    }
   ): Parameters<EntityState["write"]>[2] => {
-    const onTransaction = Context.get(entry.message.annotations, WithTransaction)
-      ? options.storage.withTransaction
-      : undefined
+    const onTransaction = !Context.get(entry.message.annotations, WithTransaction)
+      ? undefined
+      : <A, E, R>(effect: Effect.Effect<A, E, R>) =>
+        // Interrupts during commit must not hide the outcome from settle.
+        Effect.uninterruptibleMask((restore) => {
+          const transaction: RequestTransaction = {}
+          entry.transaction = transaction
+          return Effect.onExit(options.storage.withTransaction(restore(effect)), (outcome) => {
+            if (entry.transaction === transaction) entry.transaction = undefined
+            return transaction.settle ? transaction.settle(outcome) : Effect.void
+          })
+        })
     const onCaller = entry.callerScope && bindToCaller(entry.callerScope)
     if (!onCaller) return onTransaction && { onRequest: onTransaction }
     return { onRequest: onTransaction ? (effect) => onCaller(onTransaction(effect)) : onCaller }
@@ -621,13 +717,13 @@ export const make = Effect.fnUntraced(function*<
             }
           })
         }
-        activeServers.forEach((state) => {
-          if (shardId[Equal.symbol](state.address.shardId)) {
-            fibers.push(runFork(entities.removeIgnore(state.address)))
-          }
-        })
+        // Look entities up by address: `activeServers` is keyed by entity id,
+        // so an entity with the same id on another shard can hide this one.
         serverCloseLatches.forEach((latches, address) => {
           if (shardId[Equal.symbol](address.shardId)) {
+            if (!latches.closing) {
+              fibers.push(runFork(entities.removeIgnore(address)))
+            }
             fibers.push(runFork(latches.closed.await))
           }
         })

package/src/encoding/Sse.ts

@@ -92,6 +92,12 @@ export interface DecodeOptions {
    * Maximum number of string code units retained for a pending event. The default is 10 MiB.
    */
   readonly maxEventSize?: number | undefined
+  /**
+   * Also dispatch an event that sets an `event:` name but has no `data:`
+   * line, with empty `data`. The WHATWG algorithm drops such events, but some
+   * protocols use them, such as graphql-sse's `complete`. Defaults to `false`.
+   */
+  readonly dispatchEmptyNamedEvents?: boolean | undefined
 }
 
 const defaultMaxEventSize = 10 * 1024 * 1024
@@ -270,13 +276,16 @@ export const decodeDataSchema = <Type, DecodingServices, IE, Done>(
  */
 export function makeParser(onParse: (event: AnyEvent) => void, options?: DecodeOptions): Parser {
   const maxEventSize = options?.maxEventSize ?? defaultMaxEventSize
+  const dispatchEmptyNamedEvents = options?.dispatchEmptyNamedEvents === true
 
   // Processing state
   let isFirstChunk: boolean
-  let buffer: string
-  let startingPosition: number
-  let startingFieldLength: number
+  // Pieces of the unterminated line, joined once when the line ends
+  let pending: Array<string>
+  let pendingLength: number
   let discardTrailingNewline: boolean
+  // Incremented by `reset`, so `feed` can stop if `onParse` resets the parser
+  let generation = 0
 
   // Event state
   let lastEventId: string | undefined
@@ -287,10 +296,10 @@ export function makeParser(onParse: (event: AnyEvent) => void, options?: DecodeO
   return { feed, reset }
 
   function reset(): void {
+    generation++
     isFirstChunk = true
-    buffer = ""
-    startingPosition = 0
-    startingFieldLength = -1
+    pending = []
+    pendingLength = 0
     discardTrailingNewline = false
 
     lastEventId = undefined
@@ -299,120 +308,98 @@ export function makeParser(onParse: (event: AnyEvent) => void, options?: DecodeO
   }
 
   function feed(chunk: string): SseError | undefined {
-    buffer = buffer ? buffer + chunk : chunk
-
     // Strip any UTF-8 byte order mark (BOM) at the start of the stream.
     // Note that we do not strip any non - UTF8 BOM, as eventsource streams are
     // always decoded as UTF8 as per the specification.
-    if (isFirstChunk && buffer.startsWith(BOM)) {
-      buffer = buffer.slice(BOM.length)
+    if (isFirstChunk && chunk.startsWith(BOM)) {
+      chunk = chunk.slice(BOM.length)
     }
 
     isFirstChunk = false
 
-    // Set up chunk-specific processing state
-    const length = buffer.length
+    const length = chunk.length
+    const currentGeneration = generation
     let position = 0
 
-    // Read the current buffer byte by byte
     while (position < length) {
-      // EventSource allows for carriage return + line feed, which means we
-      // need to ignore a linefeed character if the previous character was a
-      // carriage return
-      // @todo refactor to reduce nesting, consider checking previous byte?
-      // @todo but consider multiple chunks etc
+      // Skip the LF of a CRLF, whose CR may have ended the previous chunk
       if (discardTrailingNewline) {
-        if (buffer[position] === "\n") {
+        discardTrailingNewline = false
+        if (chunk.charCodeAt(position) === LF) {
           ++position
+          continue
         }
-        discardTrailingNewline = false
       }
 
-      let lineLength = -1
-      let fieldLength = startingFieldLength
-      let character: string
-
-      for (let index = position + startingPosition; lineLength < 0 && index < length; ++index) {
-        character = buffer[index]
-        if (character === ":" && fieldLength < 0) {
-          fieldLength = index - position
-        } else if (character === "\r") {
-          discardTrailingNewline = true
-          lineLength = index - position
-        } else if (character === "\n") {
-          lineLength = index - position
-        }
+      let lineEnd = position
+      while (lineEnd < length) {
+        const code = chunk.charCodeAt(lineEnd)
+        if (code === CR || code === LF) break
+        ++lineEnd
       }
 
-      if (lineLength < 0) {
-        startingPosition = length - position
-        startingFieldLength = fieldLength
+      if (lineEnd === length) {
+        pending.push(chunk.slice(position))
+        pendingLength += length - position
         break
-      } else {
-        startingPosition = 0
-        startingFieldLength = -1
       }
 
-      parseEventStreamLine(buffer, position, fieldLength, lineLength)
+      discardTrailingNewline = chunk.charCodeAt(lineEnd) === CR
 
-      position += lineLength + 1
-    }
+      let line = chunk.slice(position, lineEnd)
+      if (pending.length > 0) {
+        pending.push(line)
+        line = pending.join("")
+        pending = []
+        pendingLength = 0
+      }
+      parseLine(line)
+
+      if (generation !== currentGeneration) {
+        // `onParse` reset the parser, so drop the rest of this chunk
+        return
+      }
 
-    if (position === length) {
-      // If we consumed the entire buffer to read the event, reset the buffer
-      buffer = ""
-    } else if (position > 0) {
-      // If there are bytes left to process, set the buffer to the unprocessed
-      // portion of the buffer only
-      buffer = buffer.slice(position)
+      position = lineEnd + 1
     }
 
-    if (buffer.length + data.length > maxEventSize) {
+    if (pendingLength + data.length > maxEventSize) {
       const error = new SseError({ reason: new EventTooLarge({ maxEventSize }) })
       reset()
       return error
     }
   }
 
-  function parseEventStreamLine(
-    lineBuffer: string,
-    index: number,
-    fieldLength: number,
-    lineLength: number
-  ) {
-    if (lineLength === 0) {
-      // We reached the last line of this event
-      if (data.length > 0) {
+  function parseLine(line: string) {
+    if (line.length === 0) {
+      // End of event. Clear the state first so a throwing `onParse` cannot
+      // leak it into the next event
+      const named = eventName !== undefined && eventName !== ""
+      const event = eventName || "message"
+      const eventData = data
+      eventName = undefined
+      data = ""
+      if (eventData.length > 0 || (dispatchEmptyNamedEvents && named)) {
         onParse({
           _tag: "Event",
           id: lastEventId,
-          event: eventName || "message",
-          data: data.slice(0, -1) // remove trailing newline
+          event,
+          data: eventData.slice(0, -1) // remove trailing newline
         })
-        data = ""
       }
-      eventName = undefined
       return
     }
 
-    const noValue = fieldLength < 0
-    const field = lineBuffer.slice(index, index + (noValue ? lineLength : fieldLength))
-    let step = 0
-
-    if (noValue) {
-      step = lineLength
-    } else if (lineBuffer[index + fieldLength + 1] === " ") {
-      step = fieldLength + 2
-    } else {
-      step = fieldLength + 1
+    const colon = line.indexOf(":")
+    let field = line
+    let value = ""
+    if (colon >= 0) {
+      field = line.slice(0, colon)
+      value = line.slice(line.charCodeAt(colon + 1) === SPACE ? colon + 2 : colon + 1)
     }
 
-    const position = index + step
-    const valueLength = lineLength - step
-    const value = lineBuffer.slice(position, position + valueLength).toString()
-
     if (field === "data") {
-      data += value ? `${value}\n` : "\n"
+      data += value + "\n"
     } else if (field === "event") {
       eventName = value
     } else if (field === "id" && !value.includes("\u0000")) {
@@ -425,6 +412,9 @@ export function makeParser(onParse: (event: AnyEvent) => void, options?: DecodeO
 }
 
 const BOM = "\uFEFF"
+const LF = 10
+const CR = 13
+const SPACE = 32
 
 /**
  * Stateful Server-Sent Events parser returned by `makeParser`.