|
@@ -18,6 +18,9 @@ package cn.tycoding.langchat.common.utils;
|
|
|
|
|
|
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
|
|
|
|
|
+import java.io.IOException;
|
|
|
+import java.util.concurrent.ExecutorService;
|
|
|
+
|
|
|
/**
|
|
|
* @author tycoding
|
|
|
* @since 2024/1/30
|
|
@@ -27,38 +30,64 @@ public class StreamEmitter {
|
|
|
private final SseEmitter emitter;
|
|
|
|
|
|
public StreamEmitter() {
|
|
|
- emitter = new SseEmitter(600 * 1000L);
|
|
|
+ emitter = new SseEmitter(5 * 60 * 1000L);
|
|
|
}
|
|
|
|
|
|
- public StreamEmitter(Long timeout) {
|
|
|
- emitter = new SseEmitter(timeout);
|
|
|
+ public SseEmitter get() {
|
|
|
+ return emitter;
|
|
|
}
|
|
|
|
|
|
- public SseEmitter get() {
|
|
|
+ public SseEmitter streaming(final ExecutorService executor, Runnable func) {
|
|
|
+// ExecutorService executor = Executors.newSingleThreadExecutor();
|
|
|
+
|
|
|
+ emitter.onCompletion(() -> {
|
|
|
+ System.out.println("SseEmitter 完成");
|
|
|
+ executor.shutdownNow();
|
|
|
+ });
|
|
|
+
|
|
|
+ emitter.onError((e) -> {
|
|
|
+ System.out.println("SseEmitter 出现错误: " + e.getMessage());
|
|
|
+ executor.shutdownNow();
|
|
|
+ });
|
|
|
+
|
|
|
+ emitter.onTimeout(() -> {
|
|
|
+ System.out.println("SseEmitter 超时");
|
|
|
+ emitter.complete();
|
|
|
+ executor.shutdownNow();
|
|
|
+ });
|
|
|
+ executor.execute(() -> {
|
|
|
+ try {
|
|
|
+ func.run();
|
|
|
+ } catch (Exception e) {
|
|
|
+ System.out.println("捕获到异常: " + e.getMessage());
|
|
|
+ emitter.completeWithError(e);
|
|
|
+ Thread.currentThread().interrupt();
|
|
|
+ } finally {
|
|
|
+ if (!executor.isShutdown()) {
|
|
|
+ executor.shutdownNow();
|
|
|
+ }
|
|
|
+ }
|
|
|
+ });
|
|
|
return emitter;
|
|
|
}
|
|
|
|
|
|
public void send(Object obj) {
|
|
|
try {
|
|
|
emitter.send(obj);
|
|
|
- } catch (Exception e) {
|
|
|
+ } catch (IOException e) {
|
|
|
throw new RuntimeException(e.getMessage());
|
|
|
}
|
|
|
}
|
|
|
|
|
|
public void complete() {
|
|
|
- try {
|
|
|
- emitter.complete();
|
|
|
- } catch (Exception e) {
|
|
|
- throw new RuntimeException(e.getMessage());
|
|
|
- }
|
|
|
+ emitter.complete();
|
|
|
}
|
|
|
|
|
|
public void error(String message) {
|
|
|
try {
|
|
|
emitter.send("Error: " + message);
|
|
|
emitter.complete();
|
|
|
- } catch (Exception e) {
|
|
|
+ } catch (IOException e) {
|
|
|
throw new RuntimeException(e.getMessage());
|
|
|
}
|
|
|
}
|