Java中如何用StreamingOutput将gRPC onNext字节流返回客户端？

问题背景 #

我对InputStream/OutputStream文件操作不太熟悉，需要实现CSV生成时流式传输内容，确保端点内存高效。目前已完成gRPC服务端逐行生成CSV并发送字节的逻辑，但卡在主服务器如何将gRPC接收到的字节流式返回给原终端客户端。

已实现流程 #

• 终端客户端发起REST请求，主服务器通过gRPC调用CSV服务
• CSV服务处理请求，逐行生成CSV，通过onNext以字节形式发送：

recordObserver.onNext(Response.newBuilder().setData(ByteString.copyFrom(csvRow)).build());

• 主服务器作为gRPC客户端接收每个onNext的字节块

当前代码 #

REST端点代码 #

@GET
@Produces(MediaType.APPLICATION_OCTET_STREAM)
@Path(DOWNLOAD_CSV)
StreamingOutput downloadCsv(@PathParam("Id") UUID Id) {
  return outputStream -> {
            try {
               
                mainService.downloadCsv(Id, outputStream);
                outputStream.flush();
            } catch (Exception e) {
                throw new WebApplicationException("Error occurred while attempting to download");
            }
        };

}

gRPC客户端处理代码 #

public void downloadCsv(UUID Id, OutputStream outputStream) {
    Request request = Request.newBuilder()
            .setId(Id)
            .build();

    csvServiceStub.downloadCsv(request, new StreamObserver<Response>() {
        @Override
        public void onNext(Response value) {
            try {
                byte[] bytes = value.getBytes().toByteArray();
                outputStream.write(bytes);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
        @Override
        public void onError(Throwable t) {
        }
        @Override
        public void onCompleted() {
        }
    });
}

终端请求命令 #

curl 'https://server-address/52097f0e-5dd8-49bd-98f1-69c31a73b62a/download/csv' \
  -X 'GET' \
  -H 'authority: server-address' \
  -H 'accept: application/octet-stream' \
  -H 'accept-language: en-US,en;q=0.9,ja;q=0.8' \
  -H 'authorization: OBFUSCATED' \
  -H 'content-type: application/json' \
  -H 'cookie: __zlcmid=1BzlEPGJ1UvbeK4' \
  -H 'dnt: 1' \
  -H 'origin: server-address' \
  -H 'referer: server-address' \
  --compressed

问题现象 #

1. 调用端点后终端连接立即关闭，无报错，但未输出预期的二进制内容
2. onNext方法中outputStream.write(bytes);抛出IO异常

解决方案 #

核心问题分析 #

当前代码的关键问题是gRPC调用是异步的：downloadCsv方法发起gRPC请求后立即返回，导致StreamingOutput的lambda执行完毕，REST端点直接关闭连接。后续gRPC的onNext触发时，流已经关闭，因此写入时抛出异常。

修复步骤 #

1. 用CountDownLatch同步gRPC流的生命周期，避免REST连接提前关闭
2. 在gRPC的onCompleted/onError中触发latch，确保REST端点等待流式传输完成
3. 修正gRPC响应字段的读取逻辑（和proto定义对应），并添加流刷新确保数据即时推送
4. 补充响应头，让客户端识别为CSV文件

修改后的代码 #

REST端点代码（添加头信息）

@GET
@Produces(MediaType.APPLICATION_OCTET_STREAM)
@Path(DOWNLOAD_CSV)
public StreamingOutput downloadCsv(@PathParam("Id") UUID Id, @Context HttpServletResponse response) {
    // 设置下载头，让客户端识别为CSV文件
    response.setHeader("Content-Disposition", "attachment; filename=\"data.csv\"");
    response.setContentType("text/csv; charset=utf-8");

    return outputStream -> {
        try {
            mainService.downloadCsv(Id, outputStream);
            outputStream.flush();
        } catch (Exception e) {
            throw new WebApplicationException("下载失败: " + e.getMessage(), e);
        }
    };
}

gRPC客户端处理代码（同步等待流完成）

public void downloadCsv(UUID Id, OutputStream outputStream) throws InterruptedException {
    Request request = Request.newBuilder()
            .setId(Id)
            .build();

    CountDownLatch latch = new CountDownLatch(1);
    AtomicReference<Throwable> errorRef = new AtomicReference<>();

    csvServiceStub.downloadCsv(request, new StreamObserver<Response>() {
        @Override
        public void onNext(Response value) {
            try {
                // 对应proto中的data字段，替换为实际字段名
                byte[] bytes = value.getData().toByteArray();
                outputStream.write(bytes);
                // 即时刷新，确保数据推送到客户端
                outputStream.flush();
            } catch (IOException e) {
                errorRef.set(e);
                latch.countDown();
            }
        }

        @Override
        public void onError(Throwable t) {
            errorRef.set(t);
            latch.countDown();
        }

        @Override
        public void onCompleted() {
            latch.countDown();
        }
    });

    // 等待gRPC流处理完成
    latch.await();

    // 抛出过程中出现的异常
    Throwable error = errorRef.get();
    if (error != null) {
        if (error instanceof IOException) {
            throw (IOException) error;
        } else {
            throw new RuntimeException("gRPC请求失败", error);
        }
    }
}

额外注意事项 #

• 确认gRPC的Response proto定义中字段名是data，如果是其他名称（比如bytes），需同步修改value.getData()为对应方法
• 确保CSV服务端已正确转义特殊字符（逗号、换行等），避免生成的CSV格式错误
• 测试时可暂时去掉--compressed参数，或确认服务器支持gzip压缩，避免压缩影响流式传输

内容的提问来源于stack exchange，提问作者underachiever