@yielded/agent-storage-memory 0.1.0-beta.169 → 0.1.0-beta.171

package/dist/MemorySubscriptionStore.mjs

@@ -1,9 +1,9 @@
 import { t as __exportAll } from "./rolldown-runtime-D7D4PA-g.mjs";
 import { compareScheduleNames } from "@yielded/agent/schedule-transition";
 import { Clock, Effect, Layer, Ref, Result, Schema } from "effect";
-import { AcceptedEvent, DeliveryChange, SourcePartition, SubscriptionChange, SubscriptionDelivery, SubscriptionDeliveryKey, SubscriptionError, SubscriptionFailpoint, SubscriptionKey, SubscriptionLimits, SubscriptionName, SubscriptionRecord, SubscriptionRetentionPolicy, SubscriptionScanCursors, SubscriptionStore, subscriptionDeliveryKeyString, subscriptionKeyString } from "@yielded/agent/subscription";
+import { AcceptedEvent, DeliveryChange, SourcePartition, SubscriptionChange, SubscriptionDelivery, SubscriptionDeliveryKey, SubscriptionError, SubscriptionFailpoint, SubscriptionKey, SubscriptionLimits, SubscriptionName, SubscriptionRecord, SubscriptionRetentionPolicy, SubscriptionScanCursors, SubscriptionStore, defaultSubscriptionLimits, subscriptionDeliveryKeyString, subscriptionKeyString } from "@yielded/agent/subscription";
 import { Digest } from "@yielded/agent/records";
-import { applySubscriptionChange, applySubscriptionDeliveryChange, sameAcceptedEventIdentity, subscriptionCanSelect, subscriptionDeliveryCanSelect, validateEventRetention } from "@yielded/agent/subscription-transition";
+import { applySubscriptionChange, applySubscriptionDeliveryChange, sameAcceptedEventIdentity, subscriptionCanSelect, subscriptionDeliveryCanSelect, subscriptionDeliveryUsesCapacity, subscriptionUsesCapacity, validateEventRetention } from "@yielded/agent/subscription-transition";
 //#region src/MemorySubscriptionStore.ts
 var MemorySubscriptionStore_exports = /* @__PURE__ */ __exportAll({ memorySubscriptionStoreLayer: () => memorySubscriptionStoreLayer });
 const error = (reason, code) => SubscriptionError.make({
@@ -65,6 +65,41 @@ const insertKey = (keys, key) => {
 		...keys.slice(at)
 	];
 };
+const expiryUpperBound = (expiries, nowMillis) => {
+	let low = 0;
+	let high = expiries.length;
+	while (low < high) {
+		const middle = Math.floor((low + high) / 2);
+		const expiry = expiries[middle];
+		if (expiry !== void 0 && expiry <= nowMillis) low = middle + 1;
+		else high = middle;
+	}
+	return low;
+};
+const registrationCount = (current, nowMillis, ownerId) => {
+	const expiries = ownerId === void 0 ? current.registrationExpiries : current.ownerRegistrationExpiries.get(ownerId) ?? [];
+	return expiries.length - expiryUpperBound(expiries, nowMillis);
+};
+const registrationExpiriesAfter = (current, before, after) => {
+	const expiry = (record) => record !== void 0 && (record.state === "active" || record.state === "paused") ? record.configuration.expiresAtMillis ?? Number.POSITIVE_INFINITY : null;
+	const previous = expiry(before);
+	const next = expiry(after);
+	const replace = (expiries) => {
+		if (previous === next) return expiries;
+		const updated = [...expiries];
+		if (previous !== null) updated.splice(expiryUpperBound(updated, previous) - 1, 1);
+		if (next !== null) updated.splice(expiryUpperBound(updated, next), 0, next);
+		return updated;
+	};
+	const owners = new Map(current.ownerRegistrationExpiries);
+	const expiries = replace(owners.get(after.key.ownerId) ?? []);
+	if (expiries.length === 0) owners.delete(after.key.ownerId);
+	else owners.set(after.key.ownerId, expiries);
+	return {
+		registrationExpiries: replace(current.registrationExpiries),
+		ownerRegistrationExpiries: owners
+	};
+};
 const recoveryCountsAfter = (current, next, keys) => {
 	const counts = new Map(current.recoveryCounts);
 	for (const key of keys) {
@@ -94,10 +129,13 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 		deliveries: /* @__PURE__ */ new Map(),
 		registrationIndex: /* @__PURE__ */ new Map(),
 		candidateIndex: /* @__PURE__ */ new Map(),
-		ownerRegistrationCounts: /* @__PURE__ */ new Map(),
+		registrationExpiries: [],
+		ownerRegistrationExpiries: /* @__PURE__ */ new Map(),
 		eventIndex: /* @__PURE__ */ new Map(),
 		deliveryIndex: /* @__PURE__ */ new Map(),
 		ownerDeliveryCounts: /* @__PURE__ */ new Map(),
+		liveDeliveryCount: 0,
+		pendingEventCount: 0,
 		scanCursors: {
 			events: "",
 			deliveries: "",
@@ -112,6 +150,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 		const limits = yield* validate(SubscriptionLimits, inputLimits, "register-limits");
 		yield* requirePartition(record.key, "register-partition");
 		yield* failpoint.hit("subscription:register:before");
+		const nowMillis = yield* Clock.currentTimeMillis;
 		const result = yield* Effect.uninterruptible(Ref.modify(state, (current) => {
 			const key = subscriptionKeyString(record.key);
 			const existingText = current.registrations.get(key);
@@ -123,9 +162,8 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 			if (jsonBytes(record.configuration.context) > limits.maxContextBytes) return [Result.fail(error("capacity", "context-bytes")), current];
 			if (jsonBytes(record.configuration.parameters) > limits.maxPayloadBytes) return [Result.fail(error("capacity", "parameters-bytes")), current];
 			if (record.configuration.expiresAtMillis !== null && record.configuration.expiresAtMillis - record.createdAtMillis > limits.maxLifetimeMillis) return [Result.fail(error("capacity", "lifetime")), current];
-			if (current.registrations.size >= limits.maxRegistrations) return [Result.fail(error("capacity", "registrations")), current];
-			const ownerCount = current.ownerRegistrationCounts.get(record.key.ownerId) ?? 0;
-			if (ownerCount >= limits.maxRegistrationsPerOwner) return [Result.fail(error("capacity", "owner-registrations")), current];
+			if (subscriptionUsesCapacity(record, nowMillis) && registrationCount(current, nowMillis) >= limits.maxRegistrations) return [Result.fail(error("capacity", "registrations")), current];
+			if (subscriptionUsesCapacity(record, nowMillis) && registrationCount(current, nowMillis, record.key.ownerId) >= limits.maxRegistrationsPerOwner) return [Result.fail(error("capacity", "owner-registrations")), current];
 			const assigned = {
 				...record,
 				ordinal: current.sequence + 1
@@ -145,8 +183,6 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 			const candidateIndex = new Map(current.candidateIndex);
 			const candidateKey = candidateIndexKey(assigned);
 			candidateIndex.set(candidateKey, [...candidateIndex.get(candidateKey) ?? [], key]);
-			const ownerRegistrationCounts = new Map(current.ownerRegistrationCounts);
-			ownerRegistrationCounts.set(assigned.key.ownerId, ownerCount + 1);
 			return [Result.succeed(assigned), {
 				...current,
 				sequence: assigned.ordinal,
@@ -154,7 +190,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				registrationIndex,
 				recoveryCounts: recoveryCountsAfter(current, registrationIndex, [key]),
 				candidateIndex,
-				ownerRegistrationCounts
+				...registrationExpiriesAfter(current, void 0, assigned)
 			}];
 		})).pipe(Effect.flatMap(Effect.fromResult));
 		yield* failpoint.hit("subscription:register:after");
@@ -178,11 +214,13 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 		records.sort((a, b) => a.ordinal - b.ordinal);
 		return records.slice(0, limit);
 	});
-	const change = Effect.fnUntraced(function* (input, expectedRevision, inputChange) {
+	const change = Effect.fnUntraced(function* (input, expectedRevision, inputChange, inputLimits = defaultSubscriptionLimits) {
 		const key = yield* requireKey(input, "change-key");
 		const change = yield* validate(SubscriptionChange, inputChange, "change");
+		const limits = yield* validate(SubscriptionLimits, inputLimits, "change-limits");
 		yield* validate(Schema.Int.check(Schema.isGreaterThan(0)), expectedRevision, "revision");
 		yield* failpoint.hit("subscription:change:before");
+		const nowMillis = yield* Clock.currentTimeMillis;
 		const updated = yield* Effect.uninterruptible(Ref.modify(state, (current) => {
 			const storageKey = subscriptionKeyString(key);
 			const text = current.registrations.get(storageKey);
@@ -191,6 +229,10 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 			if (Result.isFailure(existing)) return [existing, current];
 			const updated = applySubscriptionChange(existing.success, expectedRevision, change);
 			if (Result.isFailure(updated)) return [updated, current];
+			if (!subscriptionUsesCapacity(existing.success, nowMillis) && subscriptionUsesCapacity(updated.success, nowMillis)) {
+				if (registrationCount(current, nowMillis) >= limits.maxRegistrations) return [Result.fail(error("capacity", "registrations")), current];
+				if (registrationCount(current, nowMillis, key.ownerId) >= limits.maxRegistrationsPerOwner) return [Result.fail(error("capacity", "owner-registrations")), current];
+			}
 			const revised = {
 				...updated.success,
 				ordinal: current.sequence + 1
@@ -218,6 +260,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				registrations,
 				registrationIndex,
 				candidateIndex,
+				...registrationExpiriesAfter(current, existing.success, revised),
 				recoveryCounts: recoveryCountsAfter(current, registrationIndex, [storageKey])
 			}];
 		})).pipe(Effect.flatMap(Effect.fromResult));
@@ -263,6 +306,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				...current,
 				registrations,
 				registrationIndex,
+				...registrationExpiriesAfter(current, decoded.success, cancelled),
 				recoveryCounts: recoveryCountsAfter(current, registrationIndex, [storageKey])
 			}];
 		})).pipe(Effect.flatMap(Effect.fromResult));
@@ -286,7 +330,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 			const horizon = validateEventRetention(event, limits, currentTimeMillis);
 			if (Result.isFailure(horizon)) return [Result.fail(horizon.failure), current];
 			if (jsonBytes(event.payload) > limits.maxPayloadBytes) return [Result.fail(error("capacity", "payload-bytes")), current];
-			if (current.events.size - current.tombstoneCount >= limits.maxEvents) return [Result.fail(error("capacity", "events")), current];
+			if (current.pendingEventCount >= limits.maxEvents) return [Result.fail(error("capacity", "events")), current];
 			const accepted = {
 				...event,
 				cutoff: current.sequence + 1,
@@ -308,6 +352,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				sequence: accepted.cutoff,
 				events,
 				eventIndex,
+				pendingEventCount: current.pendingEventCount + 1,
 				eventKeys: insertKey(current.eventKeys, accepted.eventId),
 				retentionHorizon: limits.retention?.replayHorizonMillis ?? current.retentionHorizon,
 				retentionDeadline: limits.retention === void 0 ? current.retentionDeadline : accepted.acceptedAtMillis
@@ -347,6 +392,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 	const select = Effect.fnUntraced(function* (inputEvent, inputDeliveries, cursor, complete, nowMillis, inputLimits) {
 		const suppliedEvent = yield* validate(AcceptedEvent, inputEvent, "select-event");
 		const deliveries = yield* validate(Schema.Array(SubscriptionDelivery), inputDeliveries, "select-deliveries");
+		if (new Set(deliveries.map((delivery) => subscriptionDeliveryKeyString(delivery.key))).size !== deliveries.length) return yield* error("validation", "duplicate-delivery");
 		const limits = yield* validate(SubscriptionLimits, inputLimits, "select-limits");
 		yield* requirePartition(suppliedEvent, "select-partition");
 		yield* failpoint.hit("subscription:select:before");
@@ -365,6 +411,10 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 			const additionIndex = [];
 			const registrationUpdates = [];
 			const owners = new Map(current.ownerDeliveryCounts);
+			let registrationCapacity = {
+				registrationExpiries: current.registrationExpiries,
+				ownerRegistrationExpiries: current.ownerRegistrationExpiries
+			};
 			for (const delivery of deliveries) {
 				const recordText = current.registrations.get(subscriptionKeyString(delivery.key.subscription));
 				if (recordText === void 0) return [Result.fail(error("not-found", "subscription")), current];
@@ -413,9 +463,10 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 						recoveryAt: null,
 						recoveryKey: null
 					}]);
+					registrationCapacity = registrationExpiriesAfter(registrationCapacity, record.success, consumed);
 				}
 			}
-			if (current.deliveries.size + additions.length > limits.maxDeliveries) return [Result.fail(error("capacity", "deliveries")), current];
+			if (current.liveDeliveryCount + additions.length > limits.maxDeliveries) return [Result.fail(error("capacity", "deliveries")), current];
 			const registrations = new Map(current.registrations);
 			for (const [key, value] of updates) registrations.set(key, value);
 			const nextDeliveries = new Map(current.deliveries);
@@ -450,9 +501,12 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				events,
 				eventIndex,
 				deliveries: nextDeliveries,
+				liveDeliveryCount: current.liveDeliveryCount + additions.length,
+				pendingEventCount: current.pendingEventCount - Number(complete),
 				deliveryKeys: additions.reduce((keys, [key]) => insertKey(keys, key), current.deliveryKeys),
 				deliveryIndex,
-				ownerDeliveryCounts: owners
+				ownerDeliveryCounts: owners,
+				...registrationCapacity
 			}];
 		})).pipe(Effect.flatMap(Effect.fromResult));
 		yield* failpoint.hit("subscription:select:after");
@@ -481,7 +535,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				return sameDeliveryIdentity(existing.success, delivery) ? [Result.void, current] : [Result.fail(error("conflict", "delivery-identity")), current];
 			}
 			if (!subscriptionCanSelect(record.success, accepted.success, effectiveNowMillis, true)) return [Result.fail(error("conflict", "catch-up-eligibility")), current];
-			if (current.deliveries.size >= limits.maxDeliveries) return [Result.fail(error("capacity", "deliveries")), current];
+			if (current.liveDeliveryCount >= limits.maxDeliveries) return [Result.fail(error("capacity", "deliveries")), current];
 			const ownerCount = current.ownerDeliveryCounts.get(record.success.key.ownerId) ?? 0;
 			if (ownerCount >= limits.maxDeliveriesPerOwner) return [Result.fail(error("capacity", "owner-deliveries")), current];
 			const encodedDelivery = encode(SubscriptionDelivery, delivery, "catch-up-delivery-encode");
@@ -523,6 +577,8 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				deliveryIndex,
 				deliveryKeys: insertKey(current.deliveryKeys, key),
 				ownerDeliveryCounts,
+				...registrationExpiriesAfter(current, record.success, consumed),
+				liveDeliveryCount: current.liveDeliveryCount + 1,
 				registrations,
 				registrationIndex,
 				recoveryCounts: recoveryCountsAfter(current, registrationIndex, [consumedKey]),
@@ -586,10 +642,11 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 		items.sort((a, b) => compareScheduleNames(subscriptionDeliveryKeyString(a.key), subscriptionDeliveryKeyString(b.key)));
 		return items.slice(0, limit);
 	});
-	const changeDelivery = Effect.fnUntraced(function* (inputKey, inputDeliveryId, inputChange) {
+	const changeDelivery = Effect.fnUntraced(function* (inputKey, inputDeliveryId, inputChange, inputLimits = defaultSubscriptionLimits) {
 		const key = yield* validate(SubscriptionDeliveryKey, inputKey, "change-delivery-key");
 		const deliveryId = yield* validate(Digest, inputDeliveryId, "change-delivery-id");
 		const change = yield* validate(DeliveryChange, inputChange, "change-delivery-change");
+		const limits = yield* validate(SubscriptionLimits, inputLimits, "change-delivery-limits");
 		yield* requirePartition(key.subscription, "change-delivery-partition");
 		yield* failpoint.hit(`subscription:delivery-${change._tag.toLowerCase()}:before`);
 		const effectiveChange = change._tag === "Prepare" ? {
@@ -610,6 +667,13 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 			if (Result.isFailure(transition)) return [Result.fail(transition.failure), current];
 			if (transition.success === decoded.success) return [Result.succeed(decoded.success), current];
 			const updated = transition.success;
+			const capacityDelta = Number(subscriptionDeliveryUsesCapacity(updated)) - Number(subscriptionDeliveryUsesCapacity(decoded.success));
+			const ownerId = key.subscription.ownerId;
+			const ownerCount = current.ownerDeliveryCounts.get(ownerId) ?? 0;
+			if (capacityDelta > 0) {
+				if (current.liveDeliveryCount >= limits.maxDeliveries) return [Result.fail(error("capacity", "deliveries")), current];
+				if (ownerCount >= limits.maxDeliveriesPerOwner) return [Result.fail(error("capacity", "owner-deliveries")), current];
+			}
 			const encoded = encode(SubscriptionDelivery, updated, "change-delivery-encode");
 			if (Result.isFailure(encoded)) return [Result.fail(encoded.failure), current];
 			const deliveries = new Map(current.deliveries);
@@ -624,10 +688,15 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				observeSettlement: updated.observeSettlement ?? false,
 				nextAttemptAtMillis: updated.retry.nextAttemptAtMillis
 			});
+			const ownerDeliveryCounts = new Map(current.ownerDeliveryCounts);
+			if (ownerCount + capacityDelta === 0) ownerDeliveryCounts.delete(ownerId);
+			else ownerDeliveryCounts.set(ownerId, ownerCount + capacityDelta);
 			return [Result.succeed(updated), {
 				...current,
 				deliveries,
-				deliveryIndex
+				deliveryIndex,
+				ownerDeliveryCounts,
+				liveDeliveryCount: current.liveDeliveryCount + capacityDelta
 			}];
 		})).pipe(Effect.flatMap(Effect.fromResult));
 		yield* failpoint.hit(`subscription:delivery-${change._tag.toLowerCase()}:after`);
@@ -717,6 +786,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 			const deliveryIndex = new Map(current.deliveryIndex);
 			const eventDeliveryCounts = new Map(current.eventDeliveryCounts);
 			const ownerDeliveryCounts = new Map(current.ownerDeliveryCounts);
+			let liveDeliveryCount = current.liveDeliveryCount;
 			const deliveryPage = current.deliveryKeys.slice(upperBound(current.deliveryKeys, current.maintenanceDeliveries), upperBound(current.deliveryKeys, current.maintenanceDeliveries) + limit);
 			const eventPage = current.eventKeys.slice(upperBound(current.eventKeys, current.maintenanceEvents), upperBound(current.eventKeys, current.maintenanceEvents) + limit);
 			const cutoff = nowMillis - policy.completedRetentionMillis;
@@ -755,8 +825,13 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				removedDeliveries.add(key);
 				deliveryIndex.delete(key);
 				eventDeliveryCounts.set(event.eventId, (eventDeliveryCounts.get(event.eventId) ?? 1) - 1);
-				const owner = delivery.key.subscription.ownerId;
-				ownerDeliveryCounts.set(owner, (ownerDeliveryCounts.get(owner) ?? 1) - 1);
+				if (subscriptionDeliveryUsesCapacity(delivery)) {
+					liveDeliveryCount--;
+					const owner = delivery.key.subscription.ownerId;
+					const count = (ownerDeliveryCounts.get(owner) ?? 1) - 1;
+					if (count === 0) ownerDeliveryCounts.delete(owner);
+					else ownerDeliveryCounts.set(owner, count);
+				}
 			}
 			let tombstones = current.tombstoneCount;
 			for (const key of eventPage) {
@@ -803,6 +878,7 @@ const makeMemorySubscriptionStore = Effect.fnUntraced(function* (ownedPartition)
 				deliveryIndex,
 				eventDeliveryCounts,
 				ownerDeliveryCounts,
+				liveDeliveryCount,
 				eventKeys: removeKeys(current.eventKeys, removedEvents),
 				deliveryKeys: removeKeys(current.deliveryKeys, removedDeliveries),
 				tombstoneCount: tombstones,