Add configurable policies for consumer credit flow (#196)

* feature: configurable policies for consumer credit flow

Author: icappello

src/client.ts

@@ -11,7 +11,7 @@ import { FilterFunc, Message, Publisher, StreamPublisher } from "./publisher"
 import { ConsumerUpdateResponse } from "./requests/consumer_update_response"
 import { CreateStreamArguments, CreateStreamRequest } from "./requests/create_stream_request"
 import { CreateSuperStreamRequest } from "./requests/create_super_stream_request"
-import { CreditRequest, CreditRequestParams } from "./requests/credit_request"
+import { CreditRequest } from "./requests/credit_request"
 import { DeclarePublisherRequest } from "./requests/declare_publisher_request"
 import { DeletePublisherRequest } from "./requests/delete_publisher_request"
 import { DeleteStreamRequest } from "./requests/delete_stream_request"
@@ -42,6 +42,7 @@ import { UnsubscribeResponse } from "./responses/unsubscribe_response"
 import { SuperStreamConsumer } from "./super_stream_consumer"
 import { MessageKeyExtractorFunction, SuperStreamPublisher } from "./super_stream_publisher"
 import { DEFAULT_FRAME_MAX, REQUIRED_MANAGEMENT_VERSION, sample } from "./util"
+import { ConsumerCreditPolicy, CreditRequestWrapper, defaultCreditPolicy } from "./consumer_credit_policy"
 
 export type ConnectionClosedListener = (hadError: boolean) => void
 
@@ -54,6 +55,13 @@ type PublisherMappedValue = {
   params: DeclarePublisherParams
   filter: FilterFunc | undefined
 }
+
+type DeliverData = {
+  messages: Message[]
+  messageFilteringSupported: boolean
+  subscriptionId: number
+  consumerId: string
+}
 export class Client {
   public readonly id: string = randomUUID()
   private consumers = new Map<string, ConsumerMappedValue>()
@@ -185,7 +193,14 @@ export class Client {
 
     const consumer = new StreamConsumer(
       handle,
-      { connection, stream: params.stream, consumerId, consumerRef: params.consumerRef, offset: params.offset },
+      {
+        connection,
+        stream: params.stream,
+        consumerId,
+        consumerRef: params.consumerRef,
+        offset: params.offset,
+        creditPolicy: params.creditPolicy,
+      },
       params.filter
     )
     connection.on("metadata_update", async (metadata) => {
@@ -466,8 +481,15 @@ export class Client {
       properties["match-unfiltered"] = `${params.filter.matchUnfiltered}`
     }
 
+    const creditPolicy = params.creditPolicy || defaultCreditPolicy
+
     const res = await connection.sendAndWait<SubscribeResponse>(
-      new SubscribeRequest({ ...params, subscriptionId: consumerId, credit: 10, properties: properties })
+      new SubscribeRequest({
+        ...params,
+        subscriptionId: consumerId,
+        credit: creditPolicy.onSubscription(),
+        properties: properties,
+      })
     )
 
     if (!res.ok) {
@@ -476,50 +498,61 @@ export class Client {
     }
   }
 
-  private askForCredit(params: CreditRequestParams, connection: Connection): Promise<void> {
-    return connection.send(new CreditRequest({ ...params }))
+  private askForCredit(subscriptionId: number, connection: Connection): CreditRequestWrapper {
+    return async (howMany: number) => {
+      return connection.send(new CreditRequest({ subscriptionId: subscriptionId, credit: howMany }))
+    }
   }
 
   private getDeliverV1Callback(connectionId: string) {
     return async (response: DeliverResponse) => {
-      const { consumer, connection } = this.consumers.get(
-        computeExtendedConsumerId(response.subscriptionId, connectionId)
-      ) ?? {
-        consumer: undefined,
-        connection: undefined,
-      }
-      if (!consumer) {
-        this.logger.error(`On deliverV1 no consumer found`)
-        return
+      const deliverData = {
+        messages: response.messages,
+        subscriptionId: response.subscriptionId,
+        consumerId: computeExtendedConsumerId(response.subscriptionId, connectionId),
+        messageFilteringSupported: false,
       }
-      this.logger.debug(`on deliverV1 -> ${consumer.consumerRef}`)
-      this.logger.debug(`response.messages.length: ${response.messages.length}`)
-      await this.askForCredit({ credit: 1, subscriptionId: response.subscriptionId }, connection)
-      response.messages.map((x) => consumer.handle(x))
+      await this.handleDelivery(deliverData)
     }
   }
 
   private getDeliverV2Callback(connectionId: string) {
     return async (response: DeliverResponseV2) => {
-      const { consumer, connection } = this.consumers.get(
-        computeExtendedConsumerId(response.subscriptionId, connectionId)
-      ) ?? {
-        consumer: undefined,
-        connection: undefined,
+      const deliverData = {
+        messages: response.messages,
+        subscriptionId: response.subscriptionId,
+        consumerId: computeExtendedConsumerId(response.subscriptionId, connectionId),
+        messageFilteringSupported: true,
       }
-      if (!consumer) {
-        this.logger.error(`On deliverV2 no consumer found`)
-        return
-      }
-      this.logger.debug(`on deliverV2 -> ${consumer.consumerRef}`)
-      this.logger.debug(`response.messages.length: ${response.messages.length}`)
-      await this.askForCredit({ credit: 1, subscriptionId: response.subscriptionId }, connection)
-      if (consumer.filter) {
-        response.messages.filter((x) => consumer.filter?.postFilterFunc(x)).map((x) => consumer.handle(x))
-        return
-      }
-      response.messages.map((x) => consumer.handle(x))
+      await this.handleDelivery(deliverData)
+    }
+  }
+
+  private handleDelivery = async (deliverData: DeliverData) => {
+    const { messages, subscriptionId, consumerId, messageFilteringSupported } = deliverData
+    const { consumer, connection } = this.consumers.get(consumerId) ?? {
+      consumer: undefined,
+      connection: undefined,
     }
+    if (!consumer) {
+      this.logger.error(`On delivery, no consumer found`)
+      return
+    }
+    this.logger.debug(`on delivery -> ${consumer.consumerRef}`)
+    this.logger.debug(`response.messages.length: ${messages.length}`)
+
+    const creditRequestWrapper = this.askForCredit(subscriptionId, connection)
+    await consumer.creditPolicy.onChunkReceived(creditRequestWrapper)
+    const messageFilter =
+      messageFilteringSupported && consumer.filter?.postFilterFunc
+        ? consumer.filter?.postFilterFunc
+        : (_msg: Message) => true
+
+    messages.map((message) => {
+      if (messageFilter(message)) consumer.handle(message)
+    })
+
+    await consumer.creditPolicy.onChunkCompleted(creditRequestWrapper)
   }
 
   private getConsumerUpdateCallback(connectionId: string) {
@@ -700,9 +733,11 @@ export interface DeclareSuperStreamPublisherParams {
   routingStrategy?: RoutingStrategy
 }
 
+export type MessageFilter = (msg: Message) => boolean
+
 export interface ConsumerFilter {
   values: string[]
-  postFilterFunc: (msg: Message) => boolean
+  postFilterFunc: MessageFilter
   matchUnfiltered: boolean
 }
 
@@ -713,6 +748,7 @@ export interface DeclareConsumerParams {
   connectionClosedListener?: ConnectionClosedListener
   singleActive?: boolean
   filter?: ConsumerFilter
+  creditPolicy?: ConsumerCreditPolicy
 }
 
 export interface DeclareSuperStreamConsumerParams {

src/consumer.ts

@@ -1,6 +1,7 @@
 import { ConsumerFilter } from "./client"
 import { ConnectionInfo, Connection } from "./connection"
 import { ConnectionPool } from "./connection_pool"
+import { ConsumerCreditPolicy, defaultCreditPolicy } from "./consumer_credit_policy"
 import { Message } from "./publisher"
 import { Offset } from "./requests/subscribe_request"
 
@@ -26,6 +27,7 @@ export class StreamConsumer implements Consumer {
   public consumerRef?: string
   public offset: Offset
   private clientLocalOffset: Offset
+  private creditsHandler: ConsumerCreditPolicy
   readonly handle: ConsumerFunc
 
   constructor(
@@ -36,6 +38,7 @@ export class StreamConsumer implements Consumer {
       consumerId: number
       consumerRef?: string
       offset: Offset
+      creditPolicy?: ConsumerCreditPolicy
     },
     readonly filter?: ConsumerFilter
   ) {
@@ -46,6 +49,7 @@ export class StreamConsumer implements Consumer {
     this.offset = params.offset
     this.clientLocalOffset = this.offset.clone()
     this.connection.incrRefCount()
+    this.creditsHandler = params.creditPolicy || defaultCreditPolicy
     this.handle = this.wrapHandle(handle, params.offset)
   }
 
@@ -109,4 +113,8 @@ export class StreamConsumer implements Consumer {
   public get extendedId(): string {
     return computeExtendedConsumerId(this.consumerId, this.connection.connectionId)
   }
+
+  public get creditPolicy() {
+    return this.creditsHandler
+  }
 }

src/consumer_credit_policy.ts

@@ -0,0 +1,51 @@
+export type CreditRequestWrapper = (howMany: number) => Promise<void>
+
+export abstract class ConsumerCreditPolicy {
+  constructor(protected readonly startFrom: number) {}
+
+  public async onChunkReceived(_requestWrapper: CreditRequestWrapper) {
+    return
+  }
+
+  public async onChunkCompleted(_requestWrapper: CreditRequestWrapper) {
+    return
+  }
+
+  public async requestCredits(requestWrapper: CreditRequestWrapper, amount: number) {
+    return requestWrapper(amount)
+  }
+
+  public onSubscription() {
+    return this.startFrom
+  }
+}
+
+class NewCreditsOnChunkReceived extends ConsumerCreditPolicy {
+  constructor(startFrom: number = 1, private readonly step: number = 1) {
+    super(startFrom)
+  }
+
+  public async onChunkReceived(requestWrapper: CreditRequestWrapper) {
+    await this.requestCredits(requestWrapper, this.step)
+  }
+
+  public onSubscription(): number {
+    return this.startFrom
+  }
+}
+
+class NewCreditsOnChunkCompleted extends ConsumerCreditPolicy {
+  constructor(startFrom: number = 1, private readonly step: number = 1) {
+    super(startFrom)
+  }
+
+  public async onChunkCompleted(requestWrapper: CreditRequestWrapper) {
+    await this.requestCredits(requestWrapper, this.step)
+  }
+}
+
+export const creditsOnChunkReceived = (startFrom: number, step: number) =>
+  new NewCreditsOnChunkReceived(startFrom, step)
+export const creditsOnChunkCompleted = (startFrom: number, step: number) =>
+  new NewCreditsOnChunkCompleted(startFrom, step)
+export const defaultCreditPolicy = creditsOnChunkReceived(2, 1)