JAVA打造长连接推理网关支持持续输出流式生成能力的实现

JAVA 长连接推理网关：流式生成能力实现

大家好，今天我们来探讨如何使用 Java 构建一个支持长连接，并具备持续输出流式生成能力的推理网关。这个网关的核心目标是能够与后端推理服务建立持久连接，并通过该连接实时地将推理结果以流的形式推送给客户端，从而提升用户体验，特别是在需要实时反馈的场景下，例如实时语音识别、实时文本生成等。

1. 需求分析与架构设计

在开始编码之前，我们需要明确需求，并设计一个清晰的架构。

需求：

• 长连接支持： 网关需要能够与客户端和后端推理服务建立持久连接，避免频繁的连接建立和断开。
• 流式生成能力： 推理结果需要能够以流的形式实时推送给客户端，而不是等待所有结果生成完毕后再发送。
• 高并发支持： 网关需要能够处理大量的并发连接请求。
• 可扩展性： 架构需要具有一定的扩展性，方便后续增加新的推理服务或功能。
• 容错性： 网关需要具备一定的容错能力，能够在后端推理服务出现故障时进行处理。

架构设计：

我们将采用基于Netty的Reactor模式来实现网关。Netty是一个高性能的异步事件驱动的网络应用框架，非常适合构建高并发的网络应用。Reactor模式可以有效地处理并发连接，提高系统的吞吐量。

整体架构如下：

+-------------------+      +-------------------+      +-------------------+
|   Client          |----->|   Netty Gateway   |----->|  Inference Server |
+-------------------+      +-------------------+      +-------------------+
                       |                      |      |                   |
                       |  - Connection Pool  |      |  - Inference Logic|
                       |  - Request Handler   |      |  - Model Loading  |
                       |  - Response Streamer|      |  - Result Streaming|
                       +----------------------+      +-------------------+

• Client: 客户端，可以是浏览器、移动应用或其他服务。
• Netty Gateway: 基于 Netty 构建的网关，负责接收客户端请求，建立长连接，并将请求转发给后端推理服务。同时，负责接收推理服务返回的流式结果，并将结果推送给客户端。
• Inference Server: 后端推理服务，负责执行推理计算，并将结果以流的形式返回给网关。

组件说明：

• Connection Pool: 管理与后端推理服务的连接，避免频繁创建和销毁连接。
• Request Handler: 处理客户端请求，并将请求转换为后端推理服务所需的格式。
• Response Streamer: 负责将后端推理服务返回的流式结果推送给客户端。
• Inference Logic: 执行具体的推理逻辑，例如文本生成、图像识别等。
• Model Loading: 加载推理模型。
• Result Streaming: 将推理结果以流的形式返回给网关。

2. 核心模块实现

接下来，我们来实现网关的核心模块。

2.1 Netty Server Bootstrap

首先，我们需要创建一个 Netty Server Bootstrap，用于启动 Netty 服务器。

import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;

public class InferenceGatewayServer {

    private final int port;

    public InferenceGatewayServer(int port) {
        this.port = port;
    }

    public void run() throws Exception {
        EventLoopGroup bossGroup = new NioEventLoopGroup();
        EventLoopGroup workerGroup = new NioEventLoopGroup();
        try {
            ServerBootstrap b = new ServerBootstrap();
            b.group(bossGroup, workerGroup)
             .channel(NioServerSocketChannel.class)
             .option(ChannelOption.SO_BACKLOG, 128)
             .handler(new LoggingHandler(LogLevel.INFO))
             .childHandler(new ChannelInitializer<SocketChannel>() {
                 @Override
                 public void initChannel(SocketChannel ch) throws Exception {
                     ch.pipeline().addLast(
                             new HttpServerCodec(),
                             new HttpObjectAggregator(65536),
                             new InferenceGatewayHandler()); // Our handler
                 }
             })
             .childOption(ChannelOption.SO_KEEPALIVE, true);

            // Bind and start to accept incoming connections.
            ChannelFuture f = b.bind(port).sync();

            // Wait until the server socket is closed.
            // In this example, this does not happen, but you can do that to gracefully
            // shut down your server.
            f.channel().closeFuture().sync();
        } finally {
            workerGroup.shutdownGracefully();
            bossGroup.shutdownGracefully();
        }
    }

    public static void main(String[] args) throws Exception {
        int port = 8080;
        new InferenceGatewayServer(port).run();
    }
}

这段代码创建了一个 Netty Server，监听指定的端口。HttpServerCodec 用于处理 HTTP 协议的编解码，HttpObjectAggregator 用于将多个 HTTP 消息聚合成一个完整的 HTTP 请求。InferenceGatewayHandler 是我们自定义的 Handler，负责处理具体的业务逻辑。

2.2 InferenceGatewayHandler

InferenceGatewayHandler 负责接收客户端请求，并将请求转发给后端推理服务。同时，它还负责接收推理服务返回的流式结果，并将结果推送给客户端。

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.*;
import io.netty.util.CharsetUtil;

import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CompletableFuture;

public class InferenceGatewayHandler extends SimpleChannelInboundHandler<FullHttpRequest> {

    private static final String INFERENCE_ENDPOINT = "/inference";
    private static final String UPSTREAM_URL = "http://localhost:8081/stream"; // Replace with your inference server URL

    @Override
    public void channelReadComplete(ChannelHandlerContext ctx) throws Exception {
        ctx.flush();
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) throws Exception {
        URI uri = new URI(req.uri());
        if (!uri.getPath().equals(INFERENCE_ENDPOINT)) {
            sendError(ctx, HttpResponseStatus.NOT_FOUND);
            return;
        }

        if (req.method() != HttpMethod.POST) {
            sendError(ctx, HttpResponseStatus.METHOD_NOT_ALLOWED);
            return;
        }

        // Extract the request body and forward it to the inference server
        ByteBuf content = req.content();
        String requestBody = content.toString(CharsetUtil.UTF_8);

        // Asynchronously call the inference service and stream the results back to the client
        CompletableFuture<Void> inferenceFuture = forwardRequestAndStreamResponse(ctx, requestBody);

        inferenceFuture.exceptionally(e -> {
            System.err.println("Error during inference: " + e.getMessage());
            sendError(ctx, HttpResponseStatus.INTERNAL_SERVER_ERROR);
            return null;
        });
    }

    private CompletableFuture<Void> forwardRequestAndStreamResponse(ChannelHandlerContext ctx, String requestBody) {
        return CompletableFuture.runAsync(() -> {
            try {
                // Simulate a streaming response from the inference server (replace with actual logic)
                //This will call a real service that sends streaming data
                simulateStreamingResponse(ctx, requestBody);

            } catch (Exception e) {
                System.err.println("Error simulating streaming response: " + e.getMessage());
                throw new RuntimeException(e); // Re-throw the exception to be handled by exceptionally()
            }
        });
    }

    private void simulateStreamingResponse(ChannelHandlerContext ctx, String requestBody) throws Exception {
        // Send the initial HTTP response header with chunked transfer encoding
        FullHttpResponse response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK);
        response.headers().set(HttpHeaderNames.TRANSFER_ENCODING, HttpHeaderValues.CHUNKED);
        response.headers().set(HttpHeaderNames.CONTENT_TYPE, "text/plain; charset=UTF-8");

        ctx.write(response);

        // Simulate sending chunks of data
        String[] chunks = {"Chunk 1: Initializing...n", "Chunk 2: Processing data...n", "Chunk 3: Generating results...n", "Chunk 4: Finalizing...n"};
        for (String chunk : chunks) {
            ByteBuf content = Unpooled.copiedBuffer(chunk, StandardCharsets.UTF_8);
            LastHttpContent lastContent = new DefaultLastHttpContent(content); //Used LastHttpContent instead of HttpContent
            ctx.writeAndFlush(lastContent);
            Thread.sleep(500); // Simulate processing time
        }

        // Send the last HTTP content to indicate the end of the stream.
       // ctx.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT).addListener(ChannelFutureListener.CLOSE); // removed closing listener
        System.out.println("Streaming simulation complete for this request.");
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        cause.printStackTrace();
        ctx.close();
    }

    private static void sendError(ChannelHandlerContext ctx, HttpResponseStatus status) {
        FullHttpResponse response = new DefaultFullHttpResponse(
                HttpVersion.HTTP_1_1, status, Unpooled.copiedBuffer("Failure: " + status + "rn", CharsetUtil.UTF_8));
        response.headers().set(HttpHeaderNames.CONTENT_TYPE, "text/plain; charset=UTF-8");

        // Close the connection as soon as the error message is sent.
        ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
    }
}

这段代码接收客户端的 POST 请求，并将请求体转发给后端推理服务。forwardRequestAndStreamResponse 方法模拟了从推理服务接收流式结果，并将结果分块发送给客户端。 simulateStreamingResponse 方法模拟了推理服务的行为。

关键点：

• Chunked Transfer Encoding: 使用 HttpHeaderNames.TRANSFER_ENCODING 和 HttpHeaderValues.CHUNKED 来启用 Chunked Transfer Encoding，允许我们以流的形式发送数据。
• LastHttpContent: 使用 LastHttpContent 来标记流的结束。
• Asynchronous Processing: 使用 CompletableFuture 来异步地处理请求，避免阻塞 Netty 的 EventLoop 线程。

2.3 模拟的后端推理服务 (Inference Server)

为了完整起见，我们提供一个简单的模拟后端推理服务，它接收请求并流式地返回结果。

import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.http.*;
import io.netty.util.CharsetUtil;

import java.nio.charset.StandardCharsets;

public class MockInferenceServer {

    private final int port;

    public MockInferenceServer(int port) {
        this.port = port;
    }

    public void run() throws Exception {
        EventLoopGroup bossGroup = new NioEventLoopGroup();
        EventLoopGroup workerGroup = new NioEventLoopGroup();
        try {
            ServerBootstrap b = new ServerBootstrap();
            b.group(bossGroup, workerGroup)
                    .channel(NioServerSocketChannel.class)
                    .childHandler(new ChannelInitializer<SocketChannel>() {
                        @Override
                        protected void initChannel(SocketChannel ch) throws Exception {
                            ch.pipeline().addLast(
                                    new HttpServerCodec(),
                                    new HttpObjectAggregator(65536),
                                    new MockInferenceHandler());
                        }
                    })
                    .option(ChannelOption.SO_BACKLOG, 128)
                    .childOption(ChannelOption.SO_KEEPALIVE, true);

            ChannelFuture f = b.bind(port).sync();

            f.channel().closeFuture().sync();
        } finally {
            workerGroup.shutdownGracefully();
            bossGroup.shutdownGracefully();
        }
    }

    public static void main(String[] args) throws Exception {
        int port = 8081;
        new MockInferenceServer(port).run();
    }

    private static class MockInferenceHandler extends SimpleChannelInboundHandler<FullHttpRequest> {

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) throws Exception {
            if (!req.uri().equals("/stream")) {
                sendError(ctx, HttpResponseStatus.NOT_FOUND);
                return;
            }

            if (req.method() != HttpMethod.POST) {
                sendError(ctx, HttpResponseStatus.METHOD_NOT_ALLOWED);
                return;
            }

            // Simulate streaming response
            sendStreamingResponse(ctx);
        }

        private void sendStreamingResponse(ChannelHandlerContext ctx) throws Exception {
            // Send initial headers
            FullHttpResponse response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK);
            response.headers().set(HttpHeaderNames.TRANSFER_ENCODING, HttpHeaderValues.CHUNKED);
            response.headers().set(HttpHeaderNames.CONTENT_TYPE, "text/plain; charset=UTF-8");

            ctx.write(response);

            // Send chunks of data
            String[] chunks = {"Inference: Starting...n", "Inference: Processing...n", "Inference: Generating...n", "Inference: Complete!n"};
            for (String chunk : chunks) {
                ByteBuf content = Unpooled.copiedBuffer(chunk, StandardCharsets.UTF_8);
                LastHttpContent lastContent = new DefaultLastHttpContent(content);
                ctx.writeAndFlush(lastContent);
                Thread.sleep(500); // Simulate processing time
            }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
            cause.printStackTrace();
            ctx.close();
        }

        private void sendError(ChannelHandlerContext ctx, HttpResponseStatus status) {
            FullHttpResponse response = new DefaultFullHttpResponse(
                    HttpVersion.HTTP_1_1, status, Unpooled.copiedBuffer("Failure: " + status + "rn", CharsetUtil.UTF_8));
            response.headers().set(HttpHeaderNames.CONTENT_TYPE, "text/plain; charset=UTF-8");
            ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
        }
    }
}

这个模拟服务监听 8081 端口，并响应 /stream 路径的 POST 请求。它使用 Chunked Transfer Encoding 将模拟的推理结果分块发送给客户端。

3. 连接池的实现 (Connection Pool)

为了提高性能，我们可以使用连接池来管理与后端推理服务的连接。这里提供一个简单的基于 Apache HttpClient 的连接池示例。

import org.apache.http.client.config.RequestConfig;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;

import java.util.concurrent.TimeUnit;

public class HttpClientPool {

    private static final PoolingHttpClientConnectionManager connectionManager;
    private static final CloseableHttpClient httpClient;

    static {
        connectionManager = new PoolingHttpClientConnectionManager();
        connectionManager.setMaxTotal(200);
        connectionManager.setDefaultMaxPerRoute(20);

        RequestConfig requestConfig = RequestConfig.custom()
                .setConnectTimeout(5000)   // Connection timeout
                .setSocketTimeout(30000)   // Socket timeout
                .build();

        httpClient = HttpClientBuilder.create()
                .setConnectionManager(connectionManager)
                .setDefaultRequestConfig(requestConfig)
                .build();

        // Idle connection monitor thread
        Thread idleConnectionMonitor = new Thread(() -> {
            try {
                while (!Thread.currentThread().isInterrupted()) {
                    synchronized (this) {
                        wait(5000);
                        // Close expired connections
                        connectionManager.closeExpiredConnections();
                        // Optionally, close connections that have been idle longer than 30 sec
                        connectionManager.closeIdleConnections(30, TimeUnit.SECONDS);
                    }
                }
            } catch (InterruptedException ex) {
                // terminate
            }
        });
        idleConnectionMonitor.start();
    }

    public static CloseableHttpClient getHttpClient() {
        return httpClient;
    }

    public static void close() {
        try {
            httpClient.close();
            connectionManager.shutdown();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

这个连接池使用了 PoolingHttpClientConnectionManager 来管理连接，并定期清理过期和空闲的连接。 getHttpClient() 方法用于获取 HttpClient 实例。

使用连接池：

在 InferenceGatewayHandler 中，我们可以使用连接池来发送请求到后端推理服务。

import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.util.EntityUtils;

// Inside forwardRequestAndStreamResponse method in InferenceGatewayHandler
CloseableHttpClient httpClient = HttpClientPool.getHttpClient();
HttpPost httpPost = new HttpPost(UPSTREAM_URL);
httpPost.setEntity(new StringEntity(requestBody));

try (org.apache.http.client.methods.CloseableHttpResponse response = httpClient.execute(httpPost)) {
    // Process the response and stream the results back to the client
    // This part requires more detailed implementation to handle streaming responses from HttpClient
    String responseBody = EntityUtils.toString(response.getEntity()); //This is not streaming
    //The streaming implementation needs to parse the response as a stream.
    System.out.println("Response from inference server: " + responseBody);
} catch (Exception e) {
   //Handle the exception
}

注意： 这段代码只是一个示例。要实现真正的流式处理，需要使用 HttpClient 的流式 API，例如 InputStreamEntity，并逐块读取后端推理服务返回的数据。

4. 错误处理和容错机制

在实际应用中，我们需要考虑错误处理和容错机制，以提高系统的稳定性。

错误处理：

• 在 InferenceGatewayHandler 中，捕获异常并发送错误响应给客户端。
• 在连接池中，处理连接超时、IO 异常等。

容错机制：

• 重试机制： 在请求后端推理服务失败时，可以进行重试。
• 熔断机制： 当后端推理服务出现大量错误时，可以暂时停止向其发送请求，防止系统雪崩。
• 降级机制： 当后端推理服务不可用时，可以返回预定义的默认结果。

5. 扩展性考虑

为了方便后续扩展，我们可以采用以下措施：

• 插件化架构： 将不同的推理服务实现为插件，方便动态加载和卸载。
• 配置中心： 使用配置中心来管理网关的配置信息，例如后端推理服务的地址、连接池的大小等。
• 服务发现： 使用服务发现机制来动态地发现后端推理服务。

6. 总结一些想法

通过上述步骤，我们构建了一个基于 Java 和 Netty 的长连接推理网关，并实现了流式生成能力。该网关能够处理高并发请求，并将推理结果实时推送给客户端，提升用户体验。 实际应用中，需要根据具体需求进行调整和优化，例如实现更完善的连接池、错误处理和容错机制，以及支持更多的推理服务。 通过模块化设计和灵活的配置，可以构建一个可扩展、高性能的推理网关。