diff --git a/go.mod b/go.mod index 9f16bab62..0a6a22e92 100644 --- a/go.mod +++ b/go.mod @@ -19,7 +19,7 @@ require ( github.com/livekit/mageutil v0.0.0-20230125210925-54e8a70427c1 github.com/livekit/mediatransportutil v0.0.0-20230326055817-ed569ca13d26 github.com/livekit/protocol v1.5.2-0.20230405103303-8c8b87686b2c - github.com/livekit/psrpc v0.2.10 + github.com/livekit/psrpc v0.2.11-0.20230405191830-d76f71512630 github.com/mackerelio/go-osstat v0.2.4 github.com/magefile/mage v1.14.0 github.com/maxbrunsfeld/counterfeiter/v6 v6.6.1 diff --git a/go.sum b/go.sum index 4cf3986b3..a63672866 100644 --- a/go.sum +++ b/go.sum @@ -237,8 +237,8 @@ github.com/livekit/mediatransportutil v0.0.0-20230326055817-ed569ca13d26 h1:QlQF github.com/livekit/mediatransportutil v0.0.0-20230326055817-ed569ca13d26/go.mod h1:eDA41kiySZoG+wy4Etsjb3w0jjLx69i/vAmSjG4bteA= github.com/livekit/protocol v1.5.2-0.20230405103303-8c8b87686b2c h1:fkBV/qBvTSZogAh+dJ75WYqrAGHOpNGIi7Z1iuulHMQ= github.com/livekit/protocol v1.5.2-0.20230405103303-8c8b87686b2c/go.mod h1:UFgAWejoO4eshaaDe2jynTdQWwSktNO+8Wx19V7bs+o= -github.com/livekit/psrpc v0.2.10 h1:Ud9GzMYkKhMbB6c3RkOqe8mlRtSqtW0StyICp1jApDM= -github.com/livekit/psrpc v0.2.10/go.mod h1:K0j8f1PgLShR7Lx80KbmwFkDH2BvOnycXGV0OSRURKc= +github.com/livekit/psrpc v0.2.11-0.20230405191830-d76f71512630 h1:Rm5KLZgQxWnTidY+H8MsAV6sk1iiFxeXqPFgSLkMing= +github.com/livekit/psrpc v0.2.11-0.20230405191830-d76f71512630/go.mod h1:K0j8f1PgLShR7Lx80KbmwFkDH2BvOnycXGV0OSRURKc= github.com/mackerelio/go-osstat v0.2.4 h1:qxGbdPkFo65PXOb/F/nhDKpF2nGmGaCFDLXoZjJTtUs= github.com/mackerelio/go-osstat v0.2.4/go.mod h1:Zy+qzGdZs3A9cuIqmgbJvwbmLQH9dJvtio5ZjJTbdlQ= github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo= diff --git a/pkg/service/signal.go b/pkg/service/signal.go index 3e9eab948..2423260d6 100644 --- a/pkg/service/signal.go +++ b/pkg/service/signal.go @@ -152,13 +152,19 @@ type relaySignalResponseSink struct { psrpc.ServerStream[*rpc.RelaySignalResponse, *rpc.RelaySignalRequest] logger logger.Logger - mu sync.Mutex - queue []*livekit.SignalResponse - writing bool + mu sync.Mutex + queue []*livekit.SignalResponse + writing bool + draining bool } func (s *relaySignalResponseSink) Close() { - s.ServerStream.Close(nil) + s.mu.Lock() + s.draining = true + if !s.writing { + s.ServerStream.Close(nil) + } + s.mu.Unlock() } func (s *relaySignalResponseSink) IsClosed() bool { @@ -173,6 +179,9 @@ func (s *relaySignalResponseSink) write() { msg = s.queue[0] s.queue = s.queue[1:] } else { + if s.draining { + s.ServerStream.Close(nil) + } s.writing = false s.mu.Unlock() return @@ -189,16 +198,17 @@ func (s *relaySignalResponseSink) write() { } func (s *relaySignalResponseSink) WriteMessage(msg proto.Message) error { - if err := s.Context().Err(); err != nil { - return err + s.mu.Lock() + defer s.mu.Unlock() + + if s.draining || s.IsClosed() { + return psrpc.ErrStreamClosed } - s.mu.Lock() s.queue = append(s.queue, msg.(*livekit.SignalResponse)) if !s.writing { s.writing = true go s.write() } - s.mu.Unlock() return nil }