Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
5.1 kB · 243 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244package pluginhost
import ( "context" "errors" "fmt" "strconv" "sync" "sync/atomic"
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi")
type streamBridge struct { next atomic.Uint64 mu sync.Mutex streams map[string]*streamBridgeStream}
const streamBridgeBufferSize = 16
var errStreamBridgeClosed = errors.New("stream is not open")
type streamBridgeStream struct { chunks chan pluginapi.ExecutorStreamChunk emits chan streamBridgeEmit closes chan streamBridgeClose closed chan struct{} finished chan struct{} abort chan struct{} closeOnce sync.Once abortOnce sync.Once}
type streamBridgeEmit struct { ctx context.Context chunk pluginapi.ExecutorStreamChunk done chan error}
type streamBridgeClose struct { errorMessage string accepted chan struct{}}
type rpcStreamEmitRequest struct { StreamID string `json:"stream_id"` Payload []byte `json:"payload,omitempty"` Error string `json:"error,omitempty"`}
type rpcStreamCloseRequest struct { StreamID string `json:"stream_id"` Error string `json:"error,omitempty"`}
func newStreamBridge() *streamBridge { return &streamBridge{streams: make(map[string]*streamBridgeStream)}}
func newStreamBridgeStream() *streamBridgeStream { stream := &streamBridgeStream{ chunks: make(chan pluginapi.ExecutorStreamChunk), emits: make(chan streamBridgeEmit), closes: make(chan streamBridgeClose), closed: make(chan struct{}), finished: make(chan struct{}), abort: make(chan struct{}), } go stream.run() return stream}
func (s *streamBridgeStream) run() { defer func() { s.markClosed() close(s.chunks) close(s.finished) }()
queue := make([]pluginapi.ExecutorStreamChunk, 0, streamBridgeBufferSize) for { var emitC <-chan streamBridgeEmit if len(queue) < streamBridgeBufferSize { emitC = s.emits } var outputC chan pluginapi.ExecutorStreamChunk var next pluginapi.ExecutorStreamChunk if len(queue) > 0 { outputC = s.chunks next = queue[0] }
select { case <-s.abort: return case request := <-s.closes: s.markClosed() close(request.accepted) if request.errorMessage != "" { queue = append(queue, pluginapi.ExecutorStreamChunk{Err: fmt.Errorf("%s", request.errorMessage)}) } for len(queue) > 0 { select { case <-s.abort: return case s.chunks <- queue[0]: queue = queue[1:] } } return case request := <-emitC: if err := request.ctx.Err(); err != nil { request.done <- err continue } queue = append(queue, request.chunk) request.done <- nil case outputC <- next: queue = queue[1:] } }}
func (s *streamBridgeStream) markClosed() { if s == nil { return } s.closeOnce.Do(func() { close(s.closed) })}
func (s *streamBridgeStream) abortStream() { if s == nil { return } s.abortOnce.Do(func() { s.markClosed() close(s.abort) })}
func (s *streamBridgeStream) emit(ctx context.Context, chunk pluginapi.ExecutorStreamChunk) error { if s == nil { return errStreamBridgeClosed } if ctx == nil { ctx = context.Background() } request := streamBridgeEmit{ ctx: ctx, chunk: chunk, done: make(chan error, 1), } select { case <-ctx.Done(): return ctx.Err() case <-s.closed: return errStreamBridgeClosed case s.emits <- request: } return <-request.done}
func (s *streamBridgeStream) close(errorMessage string) { if s == nil { return } request := streamBridgeClose{ errorMessage: errorMessage, accepted: make(chan struct{}), } select { case <-s.finished: return case s.closes <- request: } select { case <-request.accepted: case <-s.finished: }}
func (b *streamBridge) open(ctx context.Context) (string, <-chan pluginapi.ExecutorStreamChunk, func()) { if b == nil { chunks := make(chan pluginapi.ExecutorStreamChunk) close(chunks) return "", chunks, func() {} } id := strconv.FormatUint(b.next.Add(1), 10) stream := newStreamBridgeStream() b.mu.Lock() b.streams[id] = stream b.mu.Unlock() cleanup := func() { b.mu.Lock() if b.streams[id] == stream { delete(b.streams, id) } b.mu.Unlock() stream.abortStream() } if ctx != nil && ctx.Done() != nil { // Abort streams canceled before ExecuteStream can install cleanupWhenStreamDone. go func() { <-ctx.Done() cleanup() }() } return id, stream.chunks, cleanup}
func (b *streamBridge) emit(ctx context.Context, id string, chunk pluginapi.ExecutorStreamChunk) error { if b == nil || id == "" { return fmt.Errorf("stream id is required") } b.mu.Lock() stream := b.streams[id] b.mu.Unlock() if stream == nil { return fmt.Errorf("stream %s is not open", id) } if err := stream.emit(ctx, chunk); err != nil { if errors.Is(err, errStreamBridgeClosed) { return fmt.Errorf("stream %s is not open", id) } return err } return nil}
func (b *streamBridge) close(id string, errorMessage string) { if b == nil || id == "" { return } b.mu.Lock() stream := b.streams[id] delete(b.streams, id) b.mu.Unlock() if stream == nil { return } stream.close(errorMessage)}