Create queue producers in pipeline phase (#8643)

dev/tool/src/index.ts

@@ -352,10 +352,8 @@ export function devTool (
         const coreWsInfo = flattenStatus(wsInfo)
         const accountClient = getAccountClient(getToolToken())
 
-        const wsProducer = getPlatformQueue('tool', cmd.region).createProducer<QueueWorkspaceMessage>(
-          toolCtx,
-          QueueTopic.Workspace
-        )
+        const queue = getPlatformQueue('tool', cmd.region)
+        const wsProducer = queue.getProducer<QueueWorkspaceMessage>(toolCtx, QueueTopic.Workspace)
 
         await createWorkspace(
           measureCtx,
@@ -377,9 +375,8 @@ export function devTool (
         })
 
         await wsProducer.send(res.workspaceUuid, [workspaceEvents.created()])
-        await wsProducer.close()
-
-        console.log('create-workspace done')
+        await queue.shutdown()
+        console.log(queue)
       })
     })
 
@@ -435,10 +432,8 @@ export function devTool (
         const coreWsInfo = flattenStatus(wsInfo)
         const measureCtx = new MeasureMetricsContext('upgrade-workspace', {})
         const accountClient = getAccountClient(getToolToken(wsInfo.uuid))
-        const wsProducer = getPlatformQueue('tool', info.region).createProducer<QueueWorkspaceMessage>(
-          toolCtx,
-          QueueTopic.Workspace
-        )
+        const queue = getPlatformQueue('tool', info.region)
+        const wsProducer = queue.getProducer<QueueWorkspaceMessage>(toolCtx, QueueTopic.Workspace)
         await upgradeWorkspace(
           measureCtx,
           version,
@@ -464,7 +459,7 @@ export function devTool (
         console.log(metricsToString(measureCtx.metrics, 'upgrade', 60))
 
         await wsProducer.send(info.uuid, [workspaceEvents.upgraded()])
-        await wsProducer.close()
+        await queue.shutdown()
         console.log('upgrade-workspace done')
       })
     })
@@ -1144,12 +1139,10 @@ export function devTool (
             storageAdapter: workspaceStorage,
             historyFile: cmd.historyFile
           })
-          const wsProducer = getPlatformQueue('tool', ws.region).createProducer<QueueWorkspaceMessage>(
-            toolCtx,
-            QueueTopic.Workspace
-          )
+          const queue = getPlatformQueue('tool', ws.region)
+          const wsProducer = queue.getProducer<QueueWorkspaceMessage>(toolCtx, QueueTopic.Workspace)
           await wsProducer.send(ws.uuid, [workspaceEvents.fullReindex()])
-          await wsProducer.close()
+          await queue.shutdown()
           await workspaceStorage?.close()
         })
       }
@@ -2092,12 +2085,10 @@ export function devTool (
         }
 
         console.log('reindex workspace', workspace)
-        const wsProducer = getPlatformQueue('tool', ws.region).createProducer<QueueWorkspaceMessage>(
-          toolCtx,
-          QueueTopic.Workspace
-        )
+        const queue = getPlatformQueue('tool', ws.region)
+        const wsProducer = queue.getProducer<QueueWorkspaceMessage>(toolCtx, QueueTopic.Workspace)
         await wsProducer.send(ws.uuid, [workspaceEvents.fullReindex()])
-        await wsProducer.close()
+        await queue.shutdown()
         console.log('done', workspace)
       })
     })

pods/fulltext/src/__tests__/indexing.spec.ts

@@ -127,7 +127,7 @@ describe('full-text-indexing', () => {
     const queue = new TestQueue(toolCtx)
     await queue.start()
     try {
-      const txProducer = queue.queue.createProducer<Tx>(toolCtx, QueueTopic.Tx)
+      const txProducer = queue.queue.getProducer<Tx>(toolCtx, QueueTopic.Tx)
       const personId = randomUUID().toString() as PersonUuid
       const wsId: WorkspaceUuid = randomUUID().toString() as WorkspaceUuid
       const token = generateToken(personId, wsId)
@@ -184,7 +184,7 @@ describe('full-text-indexing', () => {
     const queue = new TestQueue(toolCtx)
     await queue.start()
     const { pipeline, wsIds } = await preparePipeline(toolCtx, queue.queue, false) // Do not use broadcast
-    const wsProcessor = queue.queue.createProducer<QueueWorkspaceMessage>(toolCtx, QueueTopic.Workspace)
+    const wsProcessor = queue.queue.getProducer<QueueWorkspaceMessage>(toolCtx, QueueTopic.Workspace)
     try {
       const pipelineClient = wrapPipeline(toolCtx, pipeline, wsIds)
 

pods/fulltext/src/manager.ts

@@ -51,7 +51,7 @@ export class WorkspaceManager {
       this.sysHierarchy.tx(tx)
     }
 
-    this.workspaceProducer = this.opt.queue.createProducer<QueueWorkspaceMessage>(this.ctx, QueueTopic.Workspace)
+    this.workspaceProducer = this.opt.queue.getProducer<QueueWorkspaceMessage>(this.ctx, QueueTopic.Workspace)
   }
 
   shutdownInterval: any

server-plugins/calendar-resources/src/index.ts

@@ -364,19 +364,16 @@ async function putEventToQueue (
   modifiedBy: PersonId,
   changes?: Record<string, any>
 ): Promise<void> {
-  if (control.queue === undefined) {
-    return
-  }
-  const producer = control.queue.createProducer<EventCUDMessage>(
+  if (control.queue === undefined) return
+  const producer = control.queue.getProducer<EventCUDMessage>(
     control.ctx.newChild('queue', {}),
     QueueTopic.CalendarEventCUD
   )
+
   try {
     await producer.send(control.workspace.uuid, [{ action, event, modifiedBy, changes }])
   } catch (err) {
     control.ctx.error('Could not queue calendar event', { err, action, event })
-  } finally {
-    await producer.close()
   }
 }
 

server-plugins/telegram-resources/src/index.ts

@@ -260,11 +260,7 @@ function hasAttachments (doc: ActivityMessage | undefined, hierarchy: Hierarchy)
 const telegramNotificationCacheKey = 'telegram.notification.cache'
 
 async function NotificationsHandler (txes: TxCreateDoc<InboxNotification>[], control: TriggerControl): Promise<Tx[]> {
-  const queue = control.queue
-
-  if (queue === undefined) {
-    return []
-  }
+  if (control.queue === undefined) return []
 
   const availableProviders: AvailableProvidersCache = control.contextCache.get(AvailableProvidersCacheKey) ?? new Map()
 
@@ -280,14 +276,11 @@ async function NotificationsHandler (txes: TxCreateDoc<InboxNotification>[], con
   }
 
   const result: Tx[] = []
-  const producer = queue.createProducer(control.ctx, QueueTopic.TelegramBot)
-  try {
+  const producer = control.queue.getProducer<TelegramQueueMessage>(control.ctx, QueueTopic.TelegramBot)
   for (const inboxNotification of all) {
     result.push(...(await processNotification(inboxNotification, control, producer)))
   }
-  } finally {
-    await producer.close()
-  }
+
   return result
 }
 
@@ -396,15 +389,9 @@ async function ProviderSettingsHandler (
   txes: TxCUD<NotificationProviderSetting>[],
   control: TriggerControl
 ): Promise<Tx[]> {
-  const queue = control.queue
+  if (control.queue === undefined) return []
+  const producer = control.queue.getProducer<TelegramQueueMessage>(control.ctx, QueueTopic.TelegramBot)
 
-  if (queue === undefined) {
-    return []
-  }
-
-  const producer = queue.createProducer(control.ctx, QueueTopic.TelegramBot)
-
-  try {
   for (const tx of txes) {
     if (tx._class === core.class.TxCreateDoc) {
       const createTx = tx as TxCreateDoc<NotificationProviderSetting>
@@ -433,9 +420,6 @@ async function ProviderSettingsHandler (
       }
     }
   }
-  } finally {
-    await producer.close()
-  }
 
   return []
 }

server/core/src/queue/dummyQueue.ts

@@ -22,7 +22,7 @@ class DummyQueueProducer<T> implements PlatformQueueProducer<T> {
  * A dummy implementation of PlatformQueue for testing and development
  */
 export class DummyQueue implements PlatformQueue {
-  createProducer<T>(ctx: MeasureContext, topic: QueueTopic | string): PlatformQueueProducer<T> {
+  getProducer<T>(ctx: MeasureContext, topic: QueueTopic | string): PlatformQueueProducer<T> {
     return new DummyQueueProducer<T>()
   }
 

server/core/src/queue/types.ts

@@ -35,7 +35,7 @@ export interface ConsumerControl {
 }
 
 export interface PlatformQueue {
-  createProducer: <T>(ctx: MeasureContext, topic: QueueTopic | string) => PlatformQueueProducer<T>
+  getProducer: <T>(ctx: MeasureContext, topic: QueueTopic | string) => PlatformQueueProducer<T>
 
   /**
    * Create a consumer for a topic.

server/core/src/types.ts

@@ -64,7 +64,7 @@ import { type Readable } from 'stream'
 import type { DbAdapter, DomainHelper } from './adapter'
 import type { StatisticsElement, WorkspaceStatistics } from './stats'
 import { type StorageAdapter } from './storage'
-import { type PlatformQueue } from './queue'
+import { type PlatformQueueProducer, type QueueTopic, type PlatformQueue } from './queue'
 
 export interface ServerFindOptions<T extends Doc> extends FindOptions<T> {
   domain?: Domain // Allow to find for Doc's in specified domain only.
@@ -195,6 +195,7 @@ export interface PipelineContext {
   lowLevelStorage?: LowLevelStorage
   liveQuery?: LiveQuery
   queue?: PlatformQueue
+  queueProducers?: Map<QueueTopic, PlatformQueueProducer<any>>
 
   // Entry point for derived data procvessing
   derived?: Middleware

server/kafka/src/__test__/queue.spec.ts

@@ -24,7 +24,7 @@ describe('queue', () => {
         })
       })
 
-      const producer = queue.createProducer<string>(testCtx, 'qtest')
+      const producer = queue.getProducer<string>(testCtx, 'qtest')
       for (let i = 0; i < docsCount; i++) {
         await producer.send(genId, ['msg' + i])
       }
@@ -54,7 +54,7 @@ describe('queue', () => {
         })
       })
 
-      const producer = queue.createProducer<string>(testCtx, 'test')
+      const producer = queue.getProducer<string>(testCtx, 'test')
       await producer.send(genId, ['msg'])
 
       await p

server/kafka/src/index.ts

@@ -57,7 +57,7 @@ function getKafkaTopicId (topic: QueueTopic | string, config: QueueConfig): stri
 
 class PlatformQueueImpl implements PlatformQueue {
   consumers: ConsumerHandle[] = []
-  producers: PlatformQueueProducerImpl[] = []
+  producers = new Map<QueueTopic | string, PlatformQueueProducerImpl>()
   constructor (
     private readonly kafka: Kafka,
     readonly config: QueueConfig
@@ -68,7 +68,7 @@ class PlatformQueueImpl implements PlatformQueue {
   }
 
   async shutdown (): Promise<void> {
-    for (const p of this.producers) {
+    for (const [, p] of this.producers) {
       try {
         await p.close()
       } catch (err: any) {
@@ -84,10 +84,14 @@ class PlatformQueueImpl implements PlatformQueue {
     }
   }
 
-  createProducer<T>(ctx: MeasureContext, topic: QueueTopic | string): PlatformQueueProducer<T> {
-    const result = new PlatformQueueProducerImpl(ctx, this.kafka, getKafkaTopicId(topic, this.config), this)
-    this.producers.push(result)
-    return result
+  getProducer<T>(ctx: MeasureContext, topic: QueueTopic | string): PlatformQueueProducer<T> {
+    const producer = this.producers.get(topic)
+    if (producer !== undefined && !producer.isClosed()) return producer
+
+    const created = new PlatformQueueProducerImpl(ctx, this.kafka, getKafkaTopicId(topic, this.config), this)
+    this.producers.set(topic, created)
+
+    return created
   }
 
   createConsumer<T>(
@@ -152,6 +156,8 @@ class PlatformQueueImpl implements PlatformQueue {
 class PlatformQueueProducerImpl implements PlatformQueueProducer<any> {
   txProducer: Producer
   connected: Promise<void> | undefined
+  private closed = false
+
   constructor (
     readonly ctx: MeasureContext,
     kafka: Kafka,
@@ -185,7 +191,12 @@ class PlatformQueueProducerImpl implements PlatformQueueProducer<any> {
     )
   }
 
+  isClosed (): boolean {
+    return this.closed
+  }
+
   async close (): Promise<void> {
+    this.closed = true
     await this.ctx.with('disconnect', {}, () => this.txProducer.disconnect())
   }
 }

server/middleware/src/queue.ts

@@ -34,7 +34,7 @@ export class QueueMiddleware extends BaseMiddleware {
     readonly queue: PlatformQueue
   ) {
     super(context, next)
-    this.txProducer = queue.createProducer<Tx>(ctx, QueueTopic.Tx)
+    this.txProducer = queue.getProducer<Tx>(ctx, QueueTopic.Tx)
   }
 
   static create (queue: PlatformQueue): MiddlewareCreator {
@@ -53,8 +53,4 @@ export class QueueMiddleware extends BaseMiddleware {
       this.txProducer.send(this.context.workspace.uuid, ctx.contextData.broadcast.txes)
     ])
   }
-
-  async close (): Promise<void> {
-    await this.txProducer.close()
-  }
 }

server/server/src/sessionManager.ts

@@ -133,8 +133,8 @@ export class TSessionManager implements SessionManager {
         this.handleTick()
       }, 1000 / ticksPerSecond)
     }
-    this.workspaceProducer = this.queue.createProducer(ctx.newChild('queue', {}), QueueTopic.Workspace)
-    this.usersProducer = this.queue.createProducer(ctx.newChild('queue', {}), QueueTopic.Users)
+    this.workspaceProducer = this.queue.getProducer(ctx.newChild('queue', {}), QueueTopic.Workspace)
+    this.usersProducer = this.queue.getProducer(ctx.newChild('queue', {}), QueueTopic.Users)
   }
 
   scheduleMaintenance (timeMinutes: number): void {

server/workspace-service/src/index.ts

@@ -123,7 +123,7 @@ export function serveWorkspaceAccount (
 
   let canceled = false
 
-  const wsProducer = queue.createProducer<QueueWorkspaceMessage>(measureCtx, QueueTopic.Workspace)
+  const wsProducer = queue.getProducer<QueueWorkspaceMessage>(measureCtx, QueueTopic.Workspace)
   const worker = new WorkspaceWorker(
     wsProducer,
     version,
@@ -158,7 +158,6 @@ export function serveWorkspaceAccount (
 
   const close = (): void => {
     canceled = true
-    void wsProducer.close()
     void queue.shutdown()
     onClose?.()
   }