@hotmeshio/hotmesh
Version:
Permanent-Memory Workflows & AI Agents
259 lines (258 loc) • 11.2 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.PostgresSubService = void 0;
const crypto_1 = __importDefault(require("crypto"));
const enums_1 = require("../../../../modules/enums");
const key_1 = require("../../../../modules/key");
const index_1 = require("../../index");
class PostgresSubService extends index_1.SubService {
constructor(eventClient, storeClient) {
super(eventClient, storeClient);
// Instance-level subscriptions for cleanup
this.instanceSubscriptions = new Map(); // topic -> callbackKey mapping
this.instanceId = crypto_1.default.randomUUID();
}
async init(namespace = key_1.HMNS, appId, engineId, logger) {
this.namespace = namespace;
this.logger = logger;
this.appId = appId;
this.engineId = engineId;
this.setupNotificationHandler();
}
setupNotificationHandler() {
// Check if notification handler is already set up for this client
if (PostgresSubService.clientHandlers.get(this.eventClient)) {
return;
}
// Initialize subscription map for this client if it doesn't exist
if (!PostgresSubService.clientSubscriptions.has(this.eventClient)) {
PostgresSubService.clientSubscriptions.set(this.eventClient, new Map());
}
// Set up the notification handler for this client
this.eventClient.on('notification', (msg) => {
const clientSubscriptions = PostgresSubService.clientSubscriptions.get(this.eventClient);
const callbacks = clientSubscriptions?.get(msg.channel);
if (callbacks && callbacks.size > 0) {
try {
const payload = JSON.parse(msg.payload || '{}');
// Collect callbacks first to avoid modification during iteration
const callbackArray = Array.from(callbacks.entries());
// Call all callbacks
callbackArray.forEach(([callbackKey, callback], index) => {
try {
callback(msg.channel, payload);
}
catch (err) {
this.logger?.error(`Error in subscription callback for ${msg.channel}:`, err);
}
});
}
catch (err) {
this.logger?.error(`Error parsing message for topic ${msg.channel}:`, err);
}
}
});
// Mark this client as having a notification handler
PostgresSubService.clientHandlers.set(this.eventClient, true);
}
transact() {
throw new Error('Transactions are not supported in lightweight pub/sub');
}
mintKey(type, params) {
if (!this.namespace)
throw new Error('Namespace not set');
return key_1.KeyService.mintKey(this.namespace, type, params);
}
mintSafeKey(type, params) {
const originalKey = this.mintKey(type, params);
if (originalKey.length <= 63) {
return [originalKey, originalKey];
}
const { appId = '', engineId = '' } = params;
const baseKey = `${this.namespace}:${appId}:${type}`;
const maxHashLength = 63 - baseKey.length - 1; // Reserve space for `:` delimiter
const engineIdHash = crypto_1.default
.createHash('sha256')
.update(engineId)
.digest('hex')
.substring(0, maxHashLength);
const safeKey = `${baseKey}:${engineIdHash}`;
return [originalKey, safeKey];
}
async subscribe(keyType, callback, appId, topic) {
const [originalKey, safeKey] = this.mintSafeKey(keyType, {
appId,
engineId: topic,
});
// Get or create subscription map for this client
let clientSubscriptions = PostgresSubService.clientSubscriptions.get(this.eventClient);
if (!clientSubscriptions) {
clientSubscriptions = new Map();
PostgresSubService.clientSubscriptions.set(this.eventClient, clientSubscriptions);
}
// Get or create callback array for this channel
let callbacks = clientSubscriptions.get(safeKey);
if (!callbacks) {
callbacks = new Map();
clientSubscriptions.set(safeKey, callbacks);
// Start listening to the safe topic (only once per channel across all instances)
await this.eventClient.query(`LISTEN "${safeKey}"`);
}
// Generate unique callback key to avoid overwrites
const callbackKey = `${this.instanceId}-${Date.now()}-${Math.random()}`;
// Add this callback to the list with unique key
callbacks.set(callbackKey, callback);
// Track this subscription for cleanup
this.instanceSubscriptions.set(safeKey, callbackKey);
this.logger.debug(`postgres-subscribe`, {
originalKey,
safeKey,
callbackKey,
totalCallbacks: callbacks.size,
});
}
async unsubscribe(keyType, appId, topic) {
const [originalKey, safeKey] = this.mintSafeKey(keyType, {
appId,
engineId: topic,
});
const clientSubscriptions = PostgresSubService.clientSubscriptions.get(this.eventClient);
if (!clientSubscriptions) {
return;
}
const callbacks = clientSubscriptions.get(safeKey);
const callbackKey = this.instanceSubscriptions.get(safeKey);
if (!callbacks || callbacks.size === 0 || !callbackKey) {
return;
}
// Remove callback using the tracked unique key
callbacks.delete(callbackKey);
// Remove from instance tracking
this.instanceSubscriptions.delete(safeKey);
// Stop listening to the safe topic if no more callbacks exist
if (callbacks.size === 0) {
clientSubscriptions.delete(safeKey);
// Check if client is still connected before attempting to unlisten
if (!this.eventClient._ending && !this.eventClient._ended) {
try {
await this.eventClient.query(`UNLISTEN "${safeKey}"`);
}
catch (err) {
// Silently handle errors if client was closed during operation
if (!err?.message?.includes('closed') && !err?.message?.includes('queryable')) {
this.logger?.error(`Error unlistening from ${safeKey}:`, err);
}
}
}
}
this.logger.debug(`postgres-unsubscribe`, {
originalKey,
safeKey,
callbackKey,
remainingCallbacks: callbacks.size,
});
}
/**
* Cleanup method to remove all subscriptions for this instance.
* Should be called when the SubService instance is being destroyed.
*/
async cleanup() {
const clientSubscriptions = PostgresSubService.clientSubscriptions.get(this.eventClient);
if (!clientSubscriptions) {
return;
}
for (const [safeKey, callbackKey] of this.instanceSubscriptions) {
const callbacks = clientSubscriptions.get(safeKey);
if (callbacks) {
callbacks.delete(callbackKey);
// If no more callbacks exist for this channel, stop listening
if (callbacks.size === 0) {
clientSubscriptions.delete(safeKey);
try {
await this.eventClient.query(`UNLISTEN "${safeKey}"`);
}
catch (err) {
this.logger?.error(`Error unlistening from ${safeKey}:`, err);
}
}
}
}
this.instanceSubscriptions.clear();
// If no more subscriptions exist for this client, remove it from static maps
if (clientSubscriptions.size === 0) {
PostgresSubService.clientSubscriptions.delete(this.eventClient);
PostgresSubService.clientHandlers.delete(this.eventClient);
}
}
async publish(keyType, message, appId, topic) {
// Check if client is still connected before attempting to publish
// This prevents "Client was closed and is not queryable" errors during cleanup
if (this.storeClient._ending || this.storeClient._ended) {
this.logger?.debug('postgres-publish-skipped-closed-client', { appId, topic });
return false;
}
const [originalKey, safeKey] = this.mintSafeKey(keyType, {
appId,
engineId: topic,
});
let messageToPublish = message;
let payload = JSON.stringify(message);
// PostgreSQL NOTIFY has a payload limit. If job message exceeds limit,
// send a reference message instead - subscriber will fetch via getState.
if (payload.length > enums_1.HMSH_NOTIFY_PAYLOAD_LIMIT &&
message.type === 'job' &&
message.job?.metadata) {
const { jid, tpc, app, js } = message.job.metadata;
messageToPublish = {
type: 'job',
topic: message.topic,
job: {
metadata: { jid, tpc, app, js },
data: null,
},
_ref: true,
};
payload = JSON.stringify(messageToPublish);
this.logger.debug('postgres-publish-ref', {
originalKey,
safeKey,
originalSize: JSON.stringify(message).length,
refSize: payload.length,
jid,
});
}
// Publish the message using the safe topic
try {
payload = payload.replace(/'/g, "''");
await this.storeClient.query(`NOTIFY "${safeKey}", '${payload}'`);
this.logger.debug(`postgres-publish`, { originalKey, safeKey });
return true;
}
catch (err) {
// Handle gracefully if client was closed during operation
if (err?.message?.includes('closed') || err?.message?.includes('queryable')) {
this.logger?.debug('postgres-publish-failed-closed-client', {
originalKey,
safeKey,
error: err.message
});
return false;
}
// Re-throw other errors
throw err;
}
}
async psubscribe() {
throw new Error('Pattern subscriptions are not supported in PostgreSQL');
}
async punsubscribe() {
throw new Error('Pattern subscriptions are not supported in PostgreSQL');
}
}
exports.PostgresSubService = PostgresSubService;
// Static maps to manage subscriptions across all instances sharing the same client
PostgresSubService.clientSubscriptions = new Map();
PostgresSubService.clientHandlers = new Map();