All files / packages/topic/src/writer writer.ts

90% Statements 135/150
84.4% Branches 92/109
89.47% Functions 17/19
90.6% Lines 135/149

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428                                                                    10x         51086x 51086x       1098067x   1098067x 51076x     1098067x 1098049x 1x   1098048x     18x     18x 1x       17x 17x                               51086x 51086x           51086x   51086x   51086x 51086x 51086x           51086x 51086x 51086x 51086x 51086x       51086x                             51086x       51086x         51086x 8x 2x   8x 2x   8x 4x 2x               51086x 52138x   84x         84x 84x       487x 487x     487x       487x 8x 12x 12x               487x     464x 464x   464x     11x 11x 11x     6x 6x 6x 6x 4x   6x     51086x 51086x 51086x             51086x                   51086x     51086x 51086x 51086x 51086x         51086x                                 1098079x 10x   1098069x       1098069x           1098069x 1098069x 1098069x       1098069x 2x         1098067x 1098079x   1098079x                         673x 1x   672x 1x       671x 671x 671x 671x   671x 267x     404x 404x             202x 202x 202x   202x               2x       57x   8x 1x   7x       49x 49x   49x 49x     48x 3x         102062x 25x     102037x 102037x 102062x 102062x 1x     102037x       31x 31x   1x 1x               51048x         51089x 1x             51088x       1x   51087x 1x       51086x         51089x     51089x                   7x    
import { abortable } from '@ydbjs/abortable'
import type { Driver } from '@ydbjs/core'
import { loggers } from '@ydbjs/debug'
 
import { type CompressionCodec, RAW_CODEC } from '../codec.js'
import type { TX } from '../tx.js'
import { generateProducerId } from './producer-id.js'
import {
	type WriterScope,
	ackBreakdown,
	publishAcknowledged,
	publishClosed,
	publishErrored,
	publishOpened,
	publishReconnecting,
	publishSessionStarted,
	traceFlush,
} from './diagnostics.js'
import { MAX_PAYLOAD_BYTES } from './writer-state.js'
import {
	DEFAULT_FLUSH_INTERVAL_MS,
	DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_MS,
	DEFAULT_MAX_BUFFER_BYTES,
	DEFAULT_MAX_INFLIGHT_COUNT,
	DEFAULT_RECOVERY_WINDOW_MS,
	DEFAULT_UPDATE_TOKEN_INTERVAL_MS,
	type WriterRuntime,
	createWriterRuntime,
} from './writer-runtime.js'
import type { OnAckCallback, TopicWriterOptions, WriteExtra } from './types.js'
 
// Lifecycle logging on the `ydb:topic:writer` namespace — low-frequency session /
// reconnect / terminal events, so it covers the failure paths without the
// per-message noise (that lives under `ydb:topic:writer:event`).
let dbg = loggers.topic.extend('writer')
 
// Synchronous seqNo validator, owned by the facade so write() can reject bad
// input at the call site without racing the FSM's async event queue.
class SeqNoValidator {
	#mode: 'auto' | 'manual' | null = null
	#highest = 0n
 
	// Returns the message seqNo (0n means "assign at send time" in auto mode).
	validate(userSeqNo: bigint | undefined): bigint {
		let provided = userSeqNo !== undefined
 
		if (this.#mode === null) {
			this.#mode = provided ? 'manual' : 'auto'
		}
 
		if (this.#mode === 'auto') {
			if (provided) {
				throw new Error('Cannot provide a seqNo in auto mode — omit it for all messages')
			}
			return 0n
		}
 
		Iif (!provided) {
			throw new Error('Cannot omit seqNo in manual mode — provide it for all messages')
		}
		if (userSeqNo! <= this.#highest) {
			throw new Error(
				`SeqNo must be strictly increasing: got ${userSeqNo}, highest seen ${this.#highest}`
			)
		}
		this.#highest = userSeqNo!
		return userSeqNo!
	}
}
 
// The public topic writer — construct it via createTopicWriter / createTopicTxWriter,
// which resolve the producer id, wire transaction hooks, and validate options.
//
// write() is synchronous and fire-and-forget: it buffers the message and returns.
// Invalid input (bad seqNo, over-size payload, a full buffer, or a closed/failed
// writer) throws synchronously at the call site. In auto mode the final seqNo is
// assigned only when the message reaches the wire, so it is not returned — use
// flush() (or the onAck callback) for the last acknowledged seqNo.
export class TopicWriter implements AsyncDisposable, Disposable {
	#scope: WriterScope
	#codec: CompressionCodec
	#runtime: WriterRuntime
	#validator = new SeqNoValidator()
	#flushWaiters: Array<PromiseWithResolvers<bigint>> = []
 
	// Byte budget mirrored here so write() can reject a full buffer synchronously,
	// ahead of the FSM's async mailbox. Incremented on write, decremented as the
	// FSM reports acked bytes leaving the window (writer.acknowledgments.freedBytes).
	#maxBufferBytes: bigint
	#bufferedBytes = 0n
 
	#lastError: unknown = undefined
 
	#closing = false
	#closed = false
	#closedDeferred = Promise.withResolvers<void>()
 
	#onAck: OnAckCallback | undefined
	#transactional: boolean
 
	constructor(driver: Driver, options: TopicWriterOptions) {
		this.#onAck = options.onAck
		this.#transactional = options.tx !== undefined
		this.#codec = options.codec ?? RAW_CODEC
		this.#maxBufferBytes = options.maxBufferBytes ?? DEFAULT_MAX_BUFFER_BYTES
		this.#scope = { driver: driver.identity, topic: options.topic, producer: options.producer! }
 
		// One-shot effective-config snapshot for late-joining metrics/traces
		// subscribers — built from the runtime's own defaults, correct by construction.
		publishOpened(this.#scope, {
			codec: this.#codec.codec,
			maxInflightCount: options.maxInflightCount ?? DEFAULT_MAX_INFLIGHT_COUNT,
			maxBufferBytes: this.#maxBufferBytes,
			flushIntervalMs: options.flushIntervalMs ?? DEFAULT_FLUSH_INTERVAL_MS,
			updateTokenIntervalMs:
				options.updateTokenIntervalMs ?? DEFAULT_UPDATE_TOKEN_INTERVAL_MS,
			gracefulShutdownTimeoutMs:
				options.gracefulShutdownTimeoutMs ?? DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_MS,
			recoveryWindowMs: options.recoveryWindowMs ?? DEFAULT_RECOVERY_WINDOW_MS,
			retryOnSchemeError: options.retryOnSchemeError ?? false,
			...(options.partitionId !== undefined && { partitionId: options.partitionId }),
			...(options.messageGroupId !== undefined && { messageGroupId: options.messageGroupId }),
		})
 
		this.#runtime = createWriterRuntime(driver, options)
		// Fire-and-forget drain. An internal machine error rethrows from the async
		// iterator, rejecting this promise — route it to the terminal path so the
		// facade never strands a caller (nor leaks it as an unhandled rejection).
		this.#consume().catch((error) => this.#fail(error))
 
		// Tx lifecycle is wired here — not in the factory — mirroring the reader:
		// commit flushes pending writes (close drains first), rollback or a close
		// without commit drops them.
		if (options.tx) {
			options.tx.onCommit(async (signal) => {
				await this.close(signal)
			})
			options.tx.onRollback(() => {
				this.destroy(new Error('Transaction rolled back'))
			})
			options.tx.onClose((committed) => {
				if (!committed) {
					this.destroy(new Error('Transaction closed without commit'))
				}
			})
		}
	}
 
	// Drain the FSM output stream, resolving/rejecting the promises the facade owns.
	async #consume(): Promise<void> {
		for await (let output of this.#runtime.machine) {
			switch (output.type) {
				case 'writer.session':
					dbg.log(
						'session started (id=%s, lastSeqNo=%s)',
						output.sessionId,
						output.lastSeqNo
					)
					publishSessionStarted(this.#scope, output.sessionId, output.lastSeqNo)
					break
 
				case 'writer.acknowledgments':
					// Acked bytes left the window — reclaim the budget write() checks against.
					this.#bufferedBytes -= output.freedBytes
					Iif (this.#bufferedBytes < 0n) {
						this.#bufferedBytes = 0n
					}
					publishAcknowledged(
						this.#scope,
						ackBreakdown(output.acknowledgments, output.freedBytes)
					)
					if (this.#onAck) {
						for (let [seqNo, status] of output.acknowledgments) {
							try {
								this.#onAck(seqNo, status)
							} catch (error) {
								// User callback errors are logged and ignored — a throwing
								// callback must never break the writer.
								dbg.log('onAck threw: %O', error)
							}
						}
					}
					break
 
				case 'writer.flushed':
					for (let waiter of this.#flushWaiters.splice(0)) {
						waiter.resolve(output.lastSeqNo)
					}
					break
 
				case 'writer.reconnecting':
					dbg.log('reconnecting (attempt %d): %O', output.attempt, output.error)
					publishReconnecting(this.#scope, output.attempt, output.error)
					break
 
				case 'writer.error':
					dbg.log('errored: %O', output.error)
					this.#lastError = output.error
					publishErrored(this.#scope, output.error)
					for (let waiter of this.#flushWaiters.splice(0)) {
						waiter.reject(output.error)
					}
					break
 
				case 'writer.closed':
					dbg.log('closed')
					this.#markClosed()
					break
			}
		}
 
		// A graceful shutdown emits writer.closed (which set #closed) before the
		// stream ends. If it ended without one, the machine stopped without a
		// terminal output — fail defensively rather than strand callers.
		Iif (!this.#closed) {
			this.#fail(
				this.#runtime.machine.signal.reason ?? new Error('Writer stopped unexpectedly')
			)
		}
	}
 
	// Closed-side bookkeeping shared by the graceful (writer.closed) and failure
	// (#fail) paths. Idempotent — safe to reach from both.
	#markClosed(): void {
		Iif (this.#closed) {
			return
		}
		this.#closed = true
		this.#bufferedBytes = 0n
		publishClosed(this.#scope)
		for (let waiter of this.#flushWaiters.splice(0)) {
			waiter.reject(this.#lastError ?? new Error('Writer closed before flush completed'))
		}
		// Resolved AFTER writer.error was processed (emitted first), so close()
		// sees #lastError and can reject on an unclean close.
		this.#closedDeferred.resolve()
	}
 
	// Terminal fallback when the machine stops without emitting writer.error /
	// writer.closed (an internal runtime error): record the failure, then close.
	#fail(error: unknown): void {
		if (this.#closed) {
			return
		}
		if (this.#lastError === undefined) {
			this.#lastError = error
			publishErrored(this.#scope, error)
		}
		this.#markClosed()
	}
 
	write(data: Uint8Array, extra?: WriteExtra): void {
		if (this.#closed || this.#closing) {
			throw new Error('Writer is closed — cannot write messages')
		}
		Iif (this.#lastError) {
			throw new Error('Writer has failed — cannot write messages', { cause: this.#lastError })
		}
		// Size limit applies to the uncompressed payload.
		Iif (BigInt(data.length) > MAX_PAYLOAD_BYTES) {
			throw new Error(
				`Message payload of ${data.length} bytes exceeds the ${MAX_PAYLOAD_BYTES} byte limit`
			)
		}
 
		let uncompressedSize = BigInt(data.length)
		let payload = this.#codec.compress(data)
		let bufferedSize = BigInt(payload.length)
 
		// Fail-fast cap on retained (un-acknowledged) bytes to bound memory. Checked
		// before the seqNo validator mutates, so a rejected write leaves no state behind.
		if (this.#bufferedBytes + bufferedSize > this.#maxBufferBytes) {
			throw new Error(
				`Writer buffer is full: ${this.#bufferedBytes + bufferedSize} bytes would exceed the ${this.#maxBufferBytes} byte limit`
			)
		}
 
		let seqNo = this.#validator.validate(extra?.seqNo)
		this.#bufferedBytes += bufferedSize
 
		this.#runtime.machine.dispatch({
			type: 'writer.write',
			message: {
				data: payload,
				uncompressedSize,
				seqNo,
				createdAt: extra?.createdAt ?? new Date(),
				...(extra?.metadataItems && { metadataItems: extra.metadataItems }),
			},
		})
	}
 
	async flush(signal?: AbortSignal): Promise<bigint> {
		if (this.#lastError) {
			throw this.#lastError
		}
		if (this.#closed) {
			throw new Error('Writer is closed')
		}
 
		// One span per flush covers batching + server acks + any reconnect in between.
		return traceFlush(this.#scope, async () => {
			let waiter = Promise.withResolvers<bigint>()
			this.#flushWaiters.push(waiter)
			this.#runtime.machine.dispatch({ type: 'writer.flush' })
 
			if (!signal) {
				return waiter.promise
			}
 
			try {
				return await abortable(signal, waiter.promise)
			} catch (error) {
				// Abort (or a rejected flush) settled the caller's promise. Drop our
				// waiter so it neither lingers in #flushWaiters — which would grow
				// unbounded when a long-lived signal is threaded into many flush() calls
				// — nor later rejects unhandled when the writer terminates. If the FSM
				// already removed it (error/closed), indexOf is -1 and this is a no-op.
				let index = this.#flushWaiters.indexOf(waiter)
				Eif (index !== -1) {
					this.#flushWaiters.splice(index, 1)
				}
				throw error
			}
		})
	}
 
	// Debuggers and util.inspect show the constructor name, which cannot tell a tx
	// writer apart — the tag makes it render as TopicWriter [TopicTxWriter] { ... }.
	get [Symbol.toStringTag](): string {
		return this.#transactional ? 'TopicTxWriter' : 'TopicWriter'
	}
 
	async close(signal?: AbortSignal): Promise<void> {
		if (this.#closed) {
			// A close that dropped data surfaces the failure even on a repeat call.
			if (this.#lastError) {
				throw this.#lastError
			}
			return
		}
 
		// Set synchronously so a concurrent write() is rejected rather than dropped.
		this.#closing = true
		this.#runtime.machine.dispatch({ type: 'writer.close' })
 
		let closed = this.#closedDeferred.promise
		await (signal ? abortable(signal, closed) : closed)
 
		// The graceful drain failed (errored / timed out with undelivered messages).
		if (this.#lastError) {
			throw this.#lastError
		}
	}
 
	destroy(reason?: unknown): void {
		if (this.#closed) {
			return
		}
 
		this.#closing = true
		let error = reason ?? new Error('Writer destroyed')
		this.#lastError = error
		for (let waiter of this.#flushWaiters.splice(0)) {
			waiter.reject(error)
		}
 
		this.#runtime.machine.dispatch({ type: 'writer.destroy', reason: error })
	}
 
	async [Symbol.asyncDispose](): Promise<void> {
		try {
			await this.close()
		} catch (error) {
			this.destroy(error)
			throw error
		}
	}
 
	// Synchronous disposal is a hard stop — destroy() drops un-acknowledged messages
	// immediately. Use `await using` (graceful close, drains the buffer) when delivery
	// matters; a sync `using` cannot await a drain, so it must hard-stop.
	[Symbol.dispose](): void {
		this.destroy()
	}
}
 
export function createTopicWriter(driver: Driver, options: TopicWriterOptions): TopicWriter {
	if (options.partitionId !== undefined && options.messageGroupId !== undefined) {
		throw new Error(
			'partitionId and messageGroupId are mutually exclusive — provide at most one'
		)
	}
 
	// Reject send-path-deadlocking config up front rather than stalling silently:
	// maxInflightCount < 1 gates every batch, so writes would never leave the buffer.
	if (
		options.maxInflightCount !== undefined &&
		(!Number.isInteger(options.maxInflightCount) || options.maxInflightCount < 1)
	) {
		throw new Error('maxInflightCount must be a positive integer')
	}
	if (options.maxBufferBytes !== undefined && options.maxBufferBytes < 1n) {
		throw new Error('maxBufferBytes must be a positive number of bytes')
	}
 
	// A producer id is generated when omitted (zero-config writes).
	let resolved: TopicWriterOptions = {
		...options,
		producer: options.producer ?? generateProducerId(),
	}
 
	dbg.log('creating writer for topic %s', options.topic)
 
	// The constructor wires the tx lifecycle when options.tx is set.
	return new TopicWriter(driver, resolved)
}
 
// Transaction writer: writes are tagged with the tx and the tx commit waits for
// the buffered writes to flush (rollback/close drops them).
export function createTopicTxWriter(
	tx: TX,
	driver: Driver,
	options: Omit<TopicWriterOptions, 'tx'>
): TopicWriter {
	return createTopicWriter(driver, { ...options, tx })
}