|
@@ -2,10 +2,10 @@ package handler
|
|
|
|
|
|
|
|
import (
|
|
import (
|
|
|
"context"
|
|
"context"
|
|
|
- "encoding/json"
|
|
|
|
|
"errors"
|
|
"errors"
|
|
|
"strings"
|
|
"strings"
|
|
|
|
|
|
|
|
|
|
+ context2 "git.ikuban.com/server/kratos-utils/v2/transport/http/context"
|
|
|
"git.ikuban.com/server/kratos-utils/v2/transport/middleware"
|
|
"git.ikuban.com/server/kratos-utils/v2/transport/middleware"
|
|
|
"github.com/go-kratos/kratos/v2/transport/http"
|
|
"github.com/go-kratos/kratos/v2/transport/http"
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc"
|
|
@@ -28,7 +28,8 @@ func serverStreamHandler(srv any, stream grpc.StreamDesc, option *middleware.Opt
|
|
|
|
|
|
|
|
httpCtx := ctx
|
|
httpCtx := ctx
|
|
|
h := ctx.Middleware(func(ctx context.Context, _ interface{}) (interface{}, error) {
|
|
h := ctx.Middleware(func(ctx context.Context, _ interface{}) (interface{}, error) {
|
|
|
- err := stream.Handler(srv, newServerStream(ctx, dec, httpCtx.Response()))
|
|
|
|
|
|
|
+ streamCtx := context2.NewStreamContext(ctx)
|
|
|
|
|
+ err := stream.Handler(srv, newServerStream(streamCtx, dec, httpCtx.Response()))
|
|
|
return nil, err
|
|
return nil, err
|
|
|
})
|
|
})
|
|
|
|
|
|
|
@@ -45,14 +46,14 @@ func serverStreamHandler(srv any, stream grpc.StreamDesc, option *middleware.Opt
|
|
|
var _ grpc.ServerStream = (*serverStream)(nil)
|
|
var _ grpc.ServerStream = (*serverStream)(nil)
|
|
|
|
|
|
|
|
type serverStream struct {
|
|
type serverStream struct {
|
|
|
- ctx context.Context
|
|
|
|
|
|
|
+ ctx context2.StreamContext
|
|
|
dec func(interface{}) error
|
|
dec func(interface{}) error
|
|
|
w http.ResponseWriter
|
|
w http.ResponseWriter
|
|
|
metadata metadata.MD
|
|
metadata metadata.MD
|
|
|
isSendHeader bool
|
|
isSendHeader bool
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-func newServerStream(ctx context.Context, dec func(interface{}) error, w http.ResponseWriter) grpc.ServerStream {
|
|
|
|
|
|
|
+func newServerStream(ctx context2.StreamContext, dec func(interface{}) error, w http.ResponseWriter) grpc.ServerStream {
|
|
|
s := &serverStream{
|
|
s := &serverStream{
|
|
|
ctx: ctx,
|
|
ctx: ctx,
|
|
|
dec: dec,
|
|
dec: dec,
|
|
@@ -102,30 +103,39 @@ func (s *serverStream) Context() context.Context {
|
|
|
|
|
|
|
|
func (s *serverStream) SendMsg(m interface{}) error {
|
|
func (s *serverStream) SendMsg(m interface{}) error {
|
|
|
if !s.isSendHeader {
|
|
if !s.isSendHeader {
|
|
|
- err := s.SendHeader(nil)
|
|
|
|
|
|
|
+ // 设置流式响应 headers
|
|
|
|
|
+ err := s.SendHeader(metadata.MD{
|
|
|
|
|
+ "Content-Type": []string{"text/event-stream"},
|
|
|
|
|
+ "Cache-Control": []string{"no-cache"},
|
|
|
|
|
+ "Connection": []string{"keep-alive"},
|
|
|
|
|
+ "X-Accel-Buffering": []string{"no"},
|
|
|
|
|
+ })
|
|
|
if err != nil {
|
|
if err != nil {
|
|
|
return err
|
|
return err
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
-
|
|
|
|
|
switch data := m.(type) {
|
|
switch data := m.(type) {
|
|
|
case []byte:
|
|
case []byte:
|
|
|
_, err := s.w.Write(data)
|
|
_, err := s.w.Write(data)
|
|
|
if err != nil {
|
|
if err != nil {
|
|
|
return err
|
|
return err
|
|
|
}
|
|
}
|
|
|
-
|
|
|
|
|
|
|
+ case string:
|
|
|
|
|
+ _, err := s.w.Write([]byte(data))
|
|
|
|
|
+ if err != nil {
|
|
|
|
|
+ return err
|
|
|
|
|
+ }
|
|
|
default:
|
|
default:
|
|
|
- j, err := json.Marshal(data)
|
|
|
|
|
|
|
+ // 使用自定义序列化逻辑
|
|
|
|
|
+ serialized, err := s.ctx.Serialize(data)
|
|
|
if err != nil {
|
|
if err != nil {
|
|
|
return err
|
|
return err
|
|
|
}
|
|
}
|
|
|
- _, err = s.w.Write(append(j, byte('\n')))
|
|
|
|
|
|
|
+ _, err = s.w.Write(append(serialized, byte('\n')))
|
|
|
if err != nil {
|
|
if err != nil {
|
|
|
return err
|
|
return err
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
-
|
|
|
|
|
s.w.(http.Flusher).Flush()
|
|
s.w.(http.Flusher).Flush()
|
|
|
return nil
|
|
return nil
|
|
|
}
|
|
}
|