stream: fix goroutine leak

This commit is contained in:
criyle 2024-02-06 12:33:48 +00:00
parent aa41950f89
commit 581b925450
2 changed files with 13 additions and 4 deletions

View File

@ -79,15 +79,19 @@ func Start(baseCtx context.Context, s Stream, w worker.Worker, srcPrefix []strin
if err != nil {
return fmt.Errorf("convert exec request: %w", err)
}
logger.Sugar().Debugf("request: %+v", rq)
defer func() {
closeFunc := func() {
for _, f := range streamIn {
f.Close()
}
streamIn = nil
for _, f := range streamOut {
f.Close()
}
}()
streamOut = nil
}
defer closeFunc()
logger.Sugar().Debugf("request: %+v", rq)
var wg errgroup.Group
execCtx, execCancel := context.WithCancel(baseCtx)
@ -120,6 +124,8 @@ func Start(baseCtx context.Context, s Stream, w worker.Worker, srcPrefix []strin
err = sendLoop(ctx, s, outCh, rtCh, logger)
cancel()
closeFunc()
streamOut = nil
wg.Wait()
return err
}

View File

@ -197,7 +197,10 @@ func (h *wsHandle) handleStream(c *gin.Context) {
w := &streamWrapper{ctx: ctx, conn: conn, sendCh: make(chan stream.Response)}
go w.sendLoop()
stream.Start(ctx, w, h.worker, h.srcPrefix, h.logger)
if err := stream.Start(ctx, w, h.worker, h.srcPrefix, h.logger); err != nil {
h.logger.Sugar().Debugln("stream start: ", err)
c.Error(err)
}
}
type contextMap struct {