floci-ui/packages/api/src/services/msk.ts

import {
  CreateClusterCommand,
  DeleteClusterCommand,
  GetBootstrapBrokersCommand,
  KafkaClient,
  ListClustersCommand,
  type BrokerNodeGroupInfo,
} from "@aws-sdk/client-kafka";
import { Kafka } from "kafkajs";
import { awsClients } from "../aws";

export type CreateMskClusterInput = {
  clusterName: string;
  kafkaVersion: string;
  numberOfBrokerNodes: number;
  brokerNodeGroupInfo: BrokerNodeGroupInfo;
  tags?: Record<string, string>;
};

export type CreateTopicInput = {
  topic: string;
  numPartitions?: number;
  replicationFactor?: number;
  configEntries?: { name: string; value: string }[];
};

export type ProduceMessageInput = {
  key?: string;
  value: string;
  headers?: Record<string, string>;
};

type KafkaAdmin = {
  connect(): Promise<void>;
  disconnect(): Promise<void>;
  fetchTopicMetadata(input?: { topics?: string[] }): Promise<{
    topics: { name: string; partitions: unknown[] }[];
  }>;
  createTopics(input: { topics: CreateTopicInput[] }): Promise<boolean>;
  deleteTopics(input: { topics: string[] }): Promise<void>;
  fetchTopicOffsets(topic: string): Promise<{ partition: number; offset: string }[]>;
  fetchOffsets(input: {
    groupId: string;
    topics?: string[];
  }): Promise<{ topic: string; partitions: { partition: number; offset: string }[] }[]>;
};

type KafkaProducer = {
  connect(): Promise<void>;
  disconnect(): Promise<void>;
  send(input: {
    topic: string;
    messages: { key?: string; value: string; headers?: Record<string, string> }[];
  }): Promise<{ topicName: string; partition: number; baseOffset?: string }[]>;
};

type KafkaConsumer = {
  connect(): Promise<void>;
  disconnect(): Promise<void>;
  subscribe(input: { topic: string; fromBeginning: boolean }): Promise<void>;
  run(input: {
    autoCommit: boolean;
    eachMessage: (payload: {
      partition: number;
      message: {
        key: Buffer | null;
        value: Buffer | null;
        headers?: Record<string, Buffer>;
        offset: string;
        timestamp: string;
      };
    }) => Promise<void>;
  }): Promise<void>;
  stop(): Promise<void>;
};

export type KafkaClientFactory = (brokers: string[]) => {
  admin(): KafkaAdmin;
  producer(): KafkaProducer;
  consumer(input: { groupId: string }): KafkaConsumer;
};

function createKafkaClient(brokers: string[]) {
  return new Kafka({ brokers }) as unknown as ReturnType<KafkaClientFactory>;
}

function splitBootstrapBrokers(bootstrapBrokers: string): string[] {
  return bootstrapBrokers.split(",").map((broker) => broker.trim()).filter(Boolean);
}

export function createMskService(
  client: KafkaClient = awsClients.kafka,
  kafkaClientFactory: KafkaClientFactory = createKafkaClient,
) {
  async function brokers(clusterArn: string): Promise<string[]> {
    const command = new GetBootstrapBrokersCommand({ ClusterArn: clusterArn });
    command.middlewareStack.add(
      (next) => async (args) => {
        const request = args.request as { headers: Record<string, string> };
        request.headers["x-floci-internal-client"] = "true";
        return next(args);
      },
      { step: "build", name: "flociInternalMskClient" },
    );
    const response = await client.send(command);
    const bootstrapBrokers = response.BootstrapBrokerString ?? response.BootstrapBrokerStringTls;
    if (!bootstrapBrokers) {
      throw new Error(`MSK cluster ${clusterArn} did not return bootstrap brokers`);
    }
    return splitBootstrapBrokers(bootstrapBrokers);
  }

  return {
    async listClusters() {
      const clusters = [];
      let nextToken: string | undefined;
      do {
        const response = await client.send(new ListClustersCommand({ NextToken: nextToken }));
        clusters.push(...(response.ClusterInfoList ?? []));
        nextToken = response.NextToken;
      } while (nextToken);
      return clusters;
    },

    async createCluster(input: CreateMskClusterInput) {
      const response = await client.send(new CreateClusterCommand({
        ClusterName: input.clusterName,
        KafkaVersion: input.kafkaVersion,
        NumberOfBrokerNodes: input.numberOfBrokerNodes,
        BrokerNodeGroupInfo: input.brokerNodeGroupInfo,
        Tags: input.tags,
      }));
      return response;
    },

    async deleteCluster(clusterArn: string) {
      return client.send(new DeleteClusterCommand({ ClusterArn: clusterArn }));
    },

    async getBootstrapBrokers(clusterArn: string) {
      const response = await client.send(new GetBootstrapBrokersCommand({ ClusterArn: clusterArn }));
      return response.BootstrapBrokerString ?? response.BootstrapBrokerStringTls;
    },

    async listTopics(clusterArn: string) {
      const admin = kafkaClientFactory(await brokers(clusterArn)).admin();
      await admin.connect();
      try {
        const metadata = await admin.fetchTopicMetadata();
        return metadata.topics
          .filter((topic) => !topic.name.startsWith("__"))
          .map((topic) => ({
            name: topic.name,
            partitionCount: topic.partitions.length,
          }));
      } finally {
        await admin.disconnect();
      }
    },

    async createTopic(clusterArn: string, input: CreateTopicInput) {
      const admin = kafkaClientFactory(await brokers(clusterArn)).admin();
      await admin.connect();
      try {
        await admin.createTopics({ topics: [input] });
      } finally {
        await admin.disconnect();
      }
    },

    async deleteTopic(clusterArn: string, topic: string) {
      const admin = kafkaClientFactory(await brokers(clusterArn)).admin();
      await admin.connect();
      try {
        await admin.deleteTopics({ topics: [topic] });
      } finally {
        await admin.disconnect();
      }
    },

    async produceMessage(clusterArn: string, topic: string, input: ProduceMessageInput) {
      const producer = kafkaClientFactory(await brokers(clusterArn)).producer();
      await producer.connect();
      try {
        return await producer.send({ topic, messages: [input] });
      } finally {
        await producer.disconnect();
      }
    },

    async peekMessages(clusterArn: string, topic: string, limit = 50) {
      const consumer = kafkaClientFactory(await brokers(clusterArn)).consumer({
        groupId: `floci-ui-peek-${crypto.randomUUID()}`,
      });
      const messages: { partition: number; offset: string; timestamp: string; key?: string; value?: string; headers: Record<string, string> }[] = [];
      await consumer.connect();
      try {
        await consumer.subscribe({ topic, fromBeginning: true });
        await new Promise<void>(async (resolve, reject) => {
          const timeout = setTimeout(resolve, 500);
          try {
            await consumer.run({
              autoCommit: false,
              eachMessage: async ({ partition, message }) => {
                if (messages.length >= limit) {
                  return;
                }
                messages.push({
                  partition,
                  offset: message.offset,
                  timestamp: message.timestamp,
                  key: message.key?.toString(),
                  value: message.value?.toString(),
                  headers: Object.fromEntries(
                    Object.entries(message.headers ?? {}).map(([name, value]) => [name, value.toString()]),
                  ),
                });
                if (messages.length >= limit) {
                  clearTimeout(timeout);
                  resolve();
                }
              },
            });
          } catch (error) {
            clearTimeout(timeout);
            reject(error);
          }
        });
        return messages;
      } finally {
        await consumer.stop();
        await consumer.disconnect();
      }
    },

    async consumerGroupLag(clusterArn: string, groupId: string, topic?: string) {
      const admin = kafkaClientFactory(await brokers(clusterArn)).admin();
      await admin.connect();
      let completed = false;
      try {
        const offsets = await admin.fetchOffsets({ groupId, topics: topic ? [topic] : undefined });
        const lag = await Promise.all(offsets.map(async (topicOffsets) => {
          const endOffsets = await admin.fetchTopicOffsets(topicOffsets.topic);
          const endOffsetByPartition = new Map(endOffsets.map((offset) => [offset.partition, offset.offset]));
          const partitions = topicOffsets.partitions.map((offset) => ({
            partition: offset.partition,
            currentOffset: offset.offset,
            endOffset: endOffsetByPartition.get(offset.partition) ?? "0",
            lag: Math.max(0, Number(endOffsetByPartition.get(offset.partition) ?? 0) - Math.max(0, Number(offset.offset))),
          }));
          return {
            topic: topicOffsets.topic,
            partitions,
            totalLag: partitions.reduce((total, partition) => total + partition.lag, 0),
          };
        }));
        completed = true;
        return lag;
      } finally {
        try {
          await admin.disconnect();
        } catch (ignored) {
          // KafkaJS can close the coordinator socket after returning offsets.
          if (!completed) {
            throw ignored;
          }
        }
      }
    },
  };
}

export const mskService = createMskService();