Reduce usage service readiness threshold (#6759)

Author: jdolle

.changeset/gorgeous-news-report.md

@@ -0,0 +1,5 @@
+---
+'hive': patch
+---
+
+Reduce usage service readiness threshold; Disable nagles algorithm and increase keepAlive from 60s to 180s for KafkaJS

deployment/services/usage.ts

@@ -54,7 +54,7 @@ export function deployUsage({
         readinessProbe: {
           initialDelaySeconds: 10,
           periodSeconds: 5,
-          failureThreshold: 2,
+          failureThreshold: 1,
           timeoutSeconds: 5,
           endpoint: '/_readiness',
         },

packages/services/usage/src/buffer.ts

@@ -322,6 +322,7 @@ export function createKVBuffer<T>(config: {
       logger.info('Stopping buffer');
       if (timeoutId) {
         clearTimeout(timeoutId);
+        timeoutId = null;
       }
       await send({
         scheduleNextSend: false,

packages/services/usage/src/fallback-queue.ts

@@ -57,6 +57,7 @@ export function createFallbackQueue(config: {
     stop() {
       if (timeoutId !== null) {
         clearTimeout(timeoutId);
+        timeoutId = null;
       }
 
       const limit = pLimit(10);

packages/services/usage/src/index.ts

@@ -135,11 +135,12 @@ async function main() {
     const shutdown = registerShutdown({
       logger: server.log,
       async onShutdown() {
-        server.log.info('Stopping tracing handler...');
-        await tracing?.shutdown();
-
         server.log.info('Stopping service handler...');
         await Promise.all([usage.stop(), server.close()]);
+
+        // shut down tracing last so that traces are sent till the very end
+        server.log.info('Stopping tracing handler...');
+        await tracing?.shutdown();
       },
     });
 

packages/services/usage/src/usage.ts

@@ -1,4 +1,13 @@
-import { CompressionTypes, Kafka, logLevel, Partitioners, RetryOptions } from 'kafkajs';
+import net from 'net';
+import tls from 'tls';
+import {
+  CompressionTypes,
+  ISocketFactory,
+  Kafka,
+  logLevel,
+  Partitioners,
+  RetryOptions,
+} from 'kafkajs';
 import { traceInlineSync, type ServiceLogger } from '@hive/service-common';
 import type { RawOperationMap, RawReport } from '@hive/usage-common';
 import { compress } from '@hive/usage-common';
@@ -32,7 +41,7 @@ const levelMap = {
 
 const retryOptions = {
   maxRetryTime: 30_000,
-  initialRetryTime: 500,
+  initialRetryTime: 300,
   factor: 0.2,
   multiplier: 2,
   retries: 5,
@@ -103,6 +112,25 @@ export function createUsage(config: {
 }) {
   const { logger } = config;
 
+  // Default KafkaJS socketFactory implementation with minor optimizations for Azure
+  // https://github.com/tulios/kafkajs/blob/master/src/network/socketFactory.js
+  const socketFactory: ISocketFactory = ({ host, port, ssl, onConnect }) => {
+    const socket = ssl
+      ? tls.connect(
+          Object.assign({ host, port }, !net.isIP(host) ? { servername: host } : {}, ssl),
+          onConnect,
+        )
+      : net.connect({ host, port }, onConnect);
+
+    // This is equivalent to kafka's "connections.max.idle.ms"
+    socket.setKeepAlive(true, 180_000);
+    // disable nagle's algorithm to have higher throughput since this logic
+    // is already buffering messages into large payloads
+    socket.setNoDelay(true);
+
+    return socket;
+  };
+
   const kafka = new Kafka({
     clientId: 'usage',
     brokers: [config.kafka.connection.broker],
@@ -140,10 +168,11 @@ export function createUsage(config: {
       };
     },
     // settings recommended by Azure EventHub https://docs.microsoft.com/en-us/azure/event-hubs/apache-kafka-configurations
-    requestTimeout: 60_000, //
+    requestTimeout: 30_000,
     connectionTimeout: 5_000,
     authenticationTimeout: 5_000,
     retry: retryOptions,
+    socketFactory,
   });
 
   const producer = kafka.producer({
@@ -268,10 +297,25 @@ export function createUsage(config: {
     status = newStatus;
   }
 
+  producer.on(producer.events.CONNECT, () => {
+    logger.info(`Kafka producer: connected`);
+
+    if (status === Status.Unhealthy) {
+      changeStatus(Status.Ready);
+    }
+  });
+
   producer.on(producer.events.REQUEST_TIMEOUT, () => {
     logger.info('Kafka producer: request timeout');
   });
 
+  producer.on(producer.events.DISCONNECT, () => {
+    logger.info(`Kafka producer: disconnected`);
+    if (status === Status.Ready) {
+      changeStatus(Status.Unhealthy);
+    }
+  });
+
   async function stop() {
     logger.info('Started Usage shutdown...');
 
@@ -281,7 +325,6 @@ export function createUsage(config: {
     await fallback.stop();
     logger.info(`Fallback stopped`);
     await producer.disconnect();
-    logger.info(`Producer disconnected`);
 
     logger.info('Usage stopped');
   }