SSE介绍(实现流式响应)

写在前面

本文一起来看下SSE相关内容。

1：SSE是什么

全称，server-send events，基于http协议，一次http请求，server端可以分批推送数据，不同于websocket的全双工通信，SSM单向通信,一般应用于需要返回的内容较多场景中，比如大模型问答场景，答案可能比较长，就可以使用SSM来实现流式的响应。

2：例子

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>org.example</groupId>
    <artifactId>sse-test</artifactId>
    <version>1.0-SNAPSHOT</version>

    <properties>
        <maven.compiler.source>8</maven.compiler.source>
        <maven.compiler.target>8</maven.compiler.target>
    </properties>
    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>2.6.0</version>
        <relativePath/> <!-- lookup parent from repository -->
    </parent>

    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
        <dependency>
            <groupId>com.squareup.okhttp3</groupId>
            <artifactId>okhttp</artifactId>
<!--            <version>4.10.0</version>-->
            <version>3.14.9</version>
        </dependency>
        <dependency>
            <groupId>com.squareup.okhttp3</groupId>
            <artifactId>okhttp-sse</artifactId>
<!--            <version>4.10.0</version>-->
            <version>3.14.9</version>
        </dependency>
        <dependency>
            <groupId>junit</groupId>
            <artifactId>junit</artifactId>
            <version>4.13.2</version>
            <scope>test</scope>
        </dependency>

    </dependencies>

</project>

SSM server

@RestController
@RequestMapping("/sse")
public class SseController {

    @RequestMapping("/emitter")
    public SseEmitter sse(@RequestBody String inputParameter, HttpServletResponse response) {
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        SseEmitter emitter = new SseEmitter();


        // Simulate asynchronous data retrieval from the database
        new Thread(() -> {
            try {
                // Query the database based on the input parameter and send data in batches
                for (int i = 0; i < 10; i++) {
                    String data = "Data batch " + i + " for parameter: " + inputParameter;
                    emitter.send(data);
                    Thread.sleep(1000); // Simulate delay between batches
                }

                emitter.complete(); // Complete the SSE connection
            } catch (Exception e) {
                emitter.completeWithError(e); // Handle errors
            }
        }).start();

        return emitter;
    }
}

SSM client

package sse;

import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;
import org.junit.Test;

import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public class TTT {

    CountDownLatch c = new CountDownLatch(1);
    @Test
    public void aaa() throws InterruptedException {
        String json = "{\"inputParameter\": \"1234\"}";

        stream("http://localhost:8080/sse/emitter", new HashMap<>(), json, new EventSourceListener() {
            /*@Override
            public void onClosed(@NotNull EventSource eventSource) {
                System.out.println("closed");
            }

            @Override
            public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
                System.out.println("open");
            }

            @Override
            public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
                System.out.println(data);
            }*/

            @Override
            public void onClosed(EventSource eventSource) {
                System.out.println("closed");
            }

            @Override
            public void onOpen(EventSource eventSource, Response response) {
                System.out.println("open");
            }

            @Override
            public void onEvent(EventSource eventSource, String id, String type, String data) {
                System.out.println("receive data: [ " + data + " ], time:" + new Date().getTime() / 1000);
            }
        });

        c.await();
    }
    public static final MediaType MEDIA_TYPE_JSON = MediaType.parse("application/json; charset=utf-8");

    public final static int MAX_IDLE_CONNECTIONS = 20;
    public final static long KEEP_ALIVE_DURATION = 30L;

    public final static int CONNECT_TIME_OUT = 6;

    public final static int WRITE_TIME_OUT = 10;

    public final static int READ_TIME_OUT = 40;

    private final static OkHttpClient HTTP_CLIENT = new OkHttpClient.Builder()
            .connectTimeout(CONNECT_TIME_OUT, TimeUnit.SECONDS)
            .writeTimeout(WRITE_TIME_OUT, TimeUnit.SECONDS)
            .readTimeout(READ_TIME_OUT, TimeUnit.SECONDS)
            .connectionPool(new ConnectionPool(MAX_IDLE_CONNECTIONS, KEEP_ALIVE_DURATION, TimeUnit.MINUTES))
            .build();

    public static boolean stream(String url, Map<String, String> headers, String json, EventSourceListener eventSourceListener) {
        try {
            RequestBody body = RequestBody.create(MEDIA_TYPE_JSON, json);
            Request.Builder builder = new Request.Builder();
//            buildHeader(builder, headers);
            Request request = builder.url(url).post(body).build();
            EventSource.Factory factory = EventSources.createFactory(HTTP_CLIENT);
            //创建事件
//            log.info("http stream请求,url: {},参数: {}", url, json);
            factory.newEventSource(request, eventSourceListener);
            return true;
        } catch (Exception e) {
//            log.error("http stream请求，url: {} 失败 ,参数: {}", url, json, e);
        }
        return false;
    }

}