diff --git a/mtglib/internal/doppel/conn.go b/mtglib/internal/doppel/conn.go index bd76c7f..33ea88d 100644 --- a/mtglib/internal/doppel/conn.go +++ b/mtglib/internal/doppel/conn.go @@ -16,22 +16,25 @@ type Conn struct { } type connPayload struct { - ctx context.Context - ctxCancel context.CancelCauseFunc - clock Clock - wg sync.WaitGroup - syncWriteLock sync.RWMutex - writeStream bytes.Buffer - writeCond *sync.Cond + ctx context.Context + ctxCancel context.CancelCauseFunc + clock Clock + wg sync.WaitGroup + writeStream bytes.Buffer + writtenCond sync.Cond + done bool } func (c Conn) Write(p []byte) (int, error) { - c.p.syncWriteLock.RLock() - defer c.p.syncWriteLock.RUnlock() + if len(p) == 0 { + return 0, context.Cause(c.p.ctx) + } - c.p.writeCond.L.Lock() + c.p.writtenCond.L.Lock() c.p.writeStream.Write(p) - c.p.writeCond.L.Unlock() + c.p.writtenCond.L.Unlock() + + c.p.writtenCond.Signal() return len(p), context.Cause(c.p.ctx) } @@ -43,8 +46,6 @@ func (c Conn) Start() { } func (c Conn) start() { - defer c.p.writeCond.Broadcast() - buf := [tls.MaxRecordSize]byte{} for { @@ -54,11 +55,16 @@ func (c Conn) start() { case <-c.p.clock.tick: } - c.p.writeCond.L.Lock() - n, err := c.p.writeStream.Read(buf[:c.p.clock.stats.Size()]) - c.p.writeCond.L.Unlock() + size := c.p.clock.stats.Size() - if n == 0 || err != nil { + c.p.writtenCond.L.Lock() + for c.p.writeStream.Len() == 0 && !c.p.done { + c.p.writtenCond.Wait() + } + n, _ := c.p.writeStream.Read(buf[:size]) + c.p.writtenCond.L.Unlock() + + if n == 0 { continue } @@ -66,13 +72,17 @@ func (c Conn) start() { c.p.ctxCancel(err) return } - - c.p.writeCond.Signal() } } func (c Conn) Stop() { c.p.ctxCancel(nil) + + c.p.writtenCond.L.Lock() + c.p.done = true + c.p.writtenCond.L.Unlock() + c.p.writtenCond.Broadcast() + c.p.wg.Wait() } @@ -83,7 +93,9 @@ func NewConn(ctx context.Context, conn essentials.Conn, stats *Stats) Conn { p: &connPayload{ ctx: ctx, ctxCancel: cancel, - writeCond: sync.NewCond(&sync.Mutex{}), + writtenCond: sync.Cond{ + L: &sync.Mutex{}, + }, clock: Clock{ stats: stats, tick: make(chan struct{}), diff --git a/mtglib/internal/doppel/conn_test.go b/mtglib/internal/doppel/conn_test.go index e774bce..8501469 100644 --- a/mtglib/internal/doppel/conn_test.go +++ b/mtglib/internal/doppel/conn_test.go @@ -141,6 +141,37 @@ func (suite *ConnTestSuite) TestWriteReturnsErrorAfterStop() { suite.Error(err) } +func (suite *ConnTestSuite) TestStopDoesNotDeadlockWhenStartIsWaiting() { + suite.connMock. + On("Write", mock.AnythingOfType("[]uint8")). + Return(0, nil). + Maybe() + + for range 100 { + func() { + ctx, cancel := context.WithCancel(suite.ctx) + defer cancel() + + c := NewConn(ctx, suite.connMock, &Stats{ + k: 2.0, + lambda: 0.01, + }) + + done := make(chan struct{}) + go func() { + defer close(done) + c.Stop() + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + suite.Fail("Stop() deadlocked: start() likely stuck in writtenCond.Wait()") + } + }() + } +} + func (suite *ConnTestSuite) TestStopOnUnderlyingWriteError() { suite.connMock. On("Write", mock.AnythingOfType("[]uint8")).