transport: fixed lazyDnsConn earlyReserveCallWg dead lock

IrineSistiana · Nov 6, 2023 · a281fdb · a281fdb
1 parent b16a871
commit a281fdb
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/pkg/upstream/transport/conn_lazy_dial.go b/pkg/upstream/transport/conn_lazy_dial.go
@@ -133,10 +133,15 @@ type lazyDnsConnEarlyReservedExchanger lazyDnsConn
 var _ ReservedExchanger = (*lazyDnsConnEarlyReservedExchanger)(nil)
 
 func (ote *lazyDnsConnEarlyReservedExchanger) ExchangeReserved(ctx context.Context, q []byte) (resp *[]byte, err error) {
-	defer ote.WithdrawReserved()
+	defer func() {
+		ote.mu.Lock()
+		ote.reservedQuery--
+		ote.mu.Unlock()
+	}()
 
 	select {
 	case <-ctx.Done():
+		ote.earlyReserveCallWg.Done()
 		return nil, context.Cause(ctx)
 	case <-ote.dialFinished:
 		dc, err := ote.c, ote.dialErr
@@ -153,7 +158,8 @@ func (ote *lazyDnsConnEarlyReservedExchanger) ExchangeReserved(ctx context.Conte
 }
 
 func (ote *lazyDnsConnEarlyReservedExchanger) WithdrawReserved() {
+	ote.earlyReserveCallWg.Done()
 	ote.mu.Lock()
-	defer ote.mu.Unlock()
 	ote.reservedQuery--
+	ote.mu.Unlock()
 }