Lock on close wait

kylecarbs · kylecarbs · commit 83eb9cae103e · 2022-03-07T01:06:27.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -280,7 +280,9 @@ func (s *server) run(ctx context.Context) {
 			s.run(ctx)
 			return
 		}
+		s.closeMutex.Lock()
 		s.connCloseWait.Add(1)
+		s.closeMutex.Unlock()
 		go s.handlePeerConn(ctx, conn)
 	}
 }
diff --git a/peerbroker/proxy.go b/peerbroker/proxy.go
@@ -200,7 +200,7 @@ func (p *proxyDial) onClientToServerMessage(ctx context.Context, message []byte)
 		go func() {
 			defer stream.Close()
 
-			err = p.onServerToClientMessage(streamID, stream)
+			err := p.onServerToClientMessage(streamID, stream)
 			if err != nil {
 				p.logger.Debug(ctx, "failed to accept server message", slog.Error(err))
 			}

Original file line number	Diff line number	Diff line change
`@@ -280,7 +280,9 @@ func (s *server) run(ctx context.Context) {`
`280`	`280`	`s.run(ctx)`
`281`	`281`	`return`
`282`	`282`	`}`
	`283`	`+ s.closeMutex.Lock()`
`283`	`284`	`s.connCloseWait.Add(1)`
	`285`	`+ s.closeMutex.Unlock()`
`284`	`286`	`go s.handlePeerConn(ctx, conn)`
`285`	`287`	`}`
`286`	`288`	`}`
Original file line number	Diff line number	Diff line change
`@@ -200,7 +200,7 @@ func (p *proxyDial) onClientToServerMessage(ctx context.Context, message []byte)`
`200`	`200`	`go func() {`
`201`	`201`	`defer stream.Close()`
`202`	`202`
`203`		`- err = p.onServerToClientMessage(streamID, stream)`
	`203`	`+ err := p.onServerToClientMessage(streamID, stream)`
`204`	`204`	`if err != nil {`
`205`	`205`	`p.logger.Debug(ctx, "failed to accept server message", slog.Error(err))`
`206`	`206`	`}`