health api

README.md

@@ -78,6 +78,10 @@ Docker 容器镜像位于 `git-registry.mycard.moe/3rdeye/onebot-lb`。使用时
 
 该模式是默认分流策略，推荐在有机器人交互的环境下使用。
 
+### HTTP API
+
+详见本项目的 Swagger API ，在 `/docs` 路径下。
+
 ## LICENSE
 
 MIT
\ No newline at end of file

src/app.module.ts

@@ -8,6 +8,8 @@ import { OnebotGateway } from './onebot.gateway';
 import { MessageService } from './message/message.service';
 import { ReverseWsService } from './reverse-ws/reverse-ws.service';
 import { WaitBotService } from './wait-bot/wait-bot.service';
+import { HealthService } from './health/health.service';
+import { HealthController } from './health/health.controller';
 
 @Module({
   imports: [
@@ -28,6 +30,8 @@ import { WaitBotService } from './wait-bot/wait-bot.service';
     MessageService,
     ReverseWsService,
     WaitBotService,
+    HealthService,
   ],
+  controllers: [HealthController],
 })
 export class AppModule {}

src/dto/HealthInfo.dto.ts

+import { ApiProperty } from '@nestjs/swagger';
+
+export class HealthInfoDto {
+  @ApiProperty({ description: '服务名称' })
+  name: string;
+  @ApiProperty({ description: '是否健康' })
+  healthy: boolean;
+  constructor(name: string, healthy: boolean) {
+    this.name = name;
+    this.healthy = healthy;
+  }
+}

src/dto/ReturnMessage.dto.ts

 import { ApiProperty } from '@nestjs/swagger';
 import { HttpException } from '@nestjs/common';
+import { HealthInfoDto } from './HealthInfo.dto';
 
 export interface BlankReturnMessage {
   statusCode: number;
@@ -39,3 +40,13 @@ export class ReturnMessageDto<T>
     this.data = data;
   }
 }
+
+export class HealthyReturnMessageDto extends BlankReturnMessageDto {
+  @ApiProperty({ description: '健康状态', type: HealthInfoDto })
+  data: HealthInfoDto;
+}
+
+export class HealthyArrayReturnMessageDto extends BlankReturnMessageDto {
+  @ApiProperty({ description: '健康状态', type: [HealthInfoDto] })
+  data: HealthInfoDto[];
+}

src/health/health.controller.spec.ts

+import { Test, TestingModule } from '@nestjs/testing';
+import { HealthController } from './health.controller';
+
+describe('HealthController', () => {
+  let controller: HealthController;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      controllers: [HealthController],
+    }).compile();
+
+    controller = module.get<HealthController>(HealthController);
+  });
+
+  it('should be defined', () => {
+    expect(controller).toBeDefined();
+  });
+});

src/health/health.controller.ts

+import { Controller, Get, Param } from '@nestjs/common';
+import {
+  ApiOkResponse,
+  ApiOperation,
+  ApiParam,
+  ApiTags,
+} from '@nestjs/swagger';
+import {
+  BlankReturnMessageDto,
+  HealthyArrayReturnMessageDto,
+  HealthyReturnMessageDto,
+  ReturnMessageDto,
+} from '../dto/ReturnMessage.dto';
+import { HealthService } from './health.service';
+
+@Controller('health')
+@ApiTags('health')
+export class HealthController {
+  constructor(private readonly healthService: HealthService) {}
+
+  @Get('route')
+  @ApiOperation({ summary: '全体路由后端健康状态' })
+  @ApiOkResponse({ type: HealthyArrayReturnMessageDto })
+  healthOfAllRoutes() {
+    const result = this.healthService.healthOfAllRoutes();
+    return new ReturnMessageDto(200, 'success', result);
+  }
+
+  @Get('route/:name')
+  @ApiOperation({ summary: '指定路由后端健康状态' })
+  @ApiParam({ name: 'name', description: '路由名称' })
+  @ApiOkResponse({ type: HealthyReturnMessageDto })
+  healthOfRoute(@Param('name') name: string) {
+    if (!name) {
+      throw new BlankReturnMessageDto(400, 'missing name').toException();
+    }
+    const result = this.healthService.healthOfRoute(name);
+    if (!result) {
+      throw new BlankReturnMessageDto(404, 'not found').toException();
+    }
+    return new ReturnMessageDto(200, 'success', result);
+  }
+
+  @Get('bot')
+  @ApiOperation({ summary: '全体机器人后端健康状态' })
+  @ApiOkResponse({ type: HealthyArrayReturnMessageDto })
+  healthOfAllBots() {
+    const result = this.healthService.healthOfAllBots();
+    return new ReturnMessageDto(200, 'success', result);
+  }
+
+  @Get('bot/:selfId')
+  @ApiOperation({ summary: '指定机器人后端健康状态' })
+  @ApiParam({ name: 'selfId', description: '机器人 ID' })
+  @ApiOkResponse({ type: HealthyReturnMessageDto })
+  HealthOfBot(@Param('selfId') name: string) {
+    if (!name) {
+      throw new BlankReturnMessageDto(400, 'missing bot ID').toException();
+    }
+    const result = this.healthService.healthOfBot(name);
+    if (!result) {
+      throw new BlankReturnMessageDto(404, 'not found').toException();
+    }
+    return new ReturnMessageDto(200, 'success', result);
+  }
+}

src/health/health.service.spec.ts

+import { Test, TestingModule } from '@nestjs/testing';
+import { HealthService } from './health.service';
+
+describe('HealthService', () => {
+  let service: HealthService;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [HealthService],
+    }).compile();
+
+    service = module.get<HealthService>(HealthService);
+  });
+
+  it('should be defined', () => {
+    expect(service).toBeDefined();
+  });
+});

src/health/health.service.ts

+import { Injectable } from '@nestjs/common';
+import { RouteService } from '../route/route.service';
+import { InjectContext } from 'koishi-nestjs';
+import { Context } from 'koishi';
+import { HealthInfoDto } from '../dto/HealthInfo.dto';
+
+@Injectable()
+export class HealthService {
+  constructor(
+    private readonly routeService: RouteService,
+    @InjectContext() private readonly ctx: Context,
+  ) {}
+
+  healthOfAllRoutes() {
+    return this.routeService.getAllRoutes().map((r) => r.getHealthyInfo());
+  }
+
+  healthOfRoute(name: string) {
+    return this.routeService.getRouteFromName(name)?.getHealthyInfo();
+  }
+
+  healthOfAllBots() {
+    return this.ctx.bots.map(
+      (b) => new HealthInfoDto(b.selfId, b.status === 'online'),
+    );
+  }
+
+  healthOfBot(selfId: string) {
+    const bot = this.ctx.bots.find((b) => b.selfId === selfId);
+    if (!bot) {
+      return;
+    }
+    return new HealthInfoDto(bot.selfId, bot.status === 'online');
+  }
+}

src/main.ts

@@ -16,6 +16,7 @@ async function bootstrap() {
     .setTitle('onebot-lb')
     .setDescription('OneBot 负载均衡器')
     .setVersion('1.1')
+    .addTag('health', '状态检查')
     .build();
 
   const document = SwaggerModule.createDocument(app, documentConfig);

src/route/Route.ts

@@ -4,6 +4,7 @@ import { Context, Session } from 'koishi';
 import { Random, remove } from 'koishi';
 import { createHash } from 'crypto';
 import { SendTask } from '../message/message.service';
+import { HealthInfoDto } from '../dto/HealthInfo.dto';
 
 export type BalancePolicy = 'broadcast' | 'random' | 'round-robin' | 'hash';
 
@@ -57,6 +58,12 @@ export class Route implements RouteConfig {
       }, this.heartbeat);
     }
   }
+  isHealthy() {
+    return this.connections.length > 0;
+  }
+  getHealthyInfo() {
+    return new HealthInfoDto(this.name, this.isHealthy());
+  }
   send(data: any, session: Session, allConns = this.connections) {
     if (!allConns.length) {
       this.preMessages.push({ data, session });