From a281fdb3823788af41e925bc6d48632a451f0a5b Mon Sep 17 00:00:00 2001
From: Irine Sistiana <49315432+IrineSistiana@users.noreply.github.com>
Date: Mon, 6 Nov 2023 09:00:59 +0800
Subject: [PATCH] transport: fixed lazyDnsConn earlyReserveCallWg dead lock

---
 pkg/upstream/transport/conn_lazy_dial.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/pkg/upstream/transport/conn_lazy_dial.go b/pkg/upstream/transport/conn_lazy_dial.go
index 63a7b6a41..a8c4e0be5 100644
--- 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()
 }