fix: avoid wrong lock deletation

Author: LostAttractor

control/dns_control.go

@@ -14,6 +14,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/daeuniverse/dae/common/consts"
@@ -82,6 +83,11 @@ type DnsController struct {
 	dnsForwarderCache   map[dnsForwarderKey]DnsForwarder
 }
 
+type handlingState struct {
+	mu  sync.Mutex
+	ref uint32
+}
+
 func parseIpVersionPreference(prefer int) (uint16, error) {
 	switch prefer := IpVersionPrefer(prefer); prefer {
 	case IpVersionPrefer_No:
@@ -453,11 +459,17 @@ func (c *DnsController) handle_(
 	}
 
 	// No parallel for the same lookup.
-	_mu, _ := c.handling.LoadOrStore(cacheKey, new(sync.Mutex))
-	mu := _mu.(*sync.Mutex)
-	mu.Lock()
-	defer mu.Unlock()
-	defer c.handling.Delete(cacheKey)
+	handlingState_, _ := c.handling.LoadOrStore(cacheKey, new(handlingState))
+	handlingState := handlingState_.(*handlingState)
+	atomic.AddUint32(&handlingState.ref, 1)
+	handlingState.mu.Lock()
+	defer func() {
+		handlingState.mu.Unlock()
+		atomic.AddUint32(&handlingState.ref, ^uint32(0))
+		if atomic.LoadUint32(&handlingState.ref) == 0 {
+			c.handling.Delete(cacheKey)
+		}
+	}()
 
 	if resp := c.LookupDnsRespCache_(dnsMessage, cacheKey, false); resp != nil {
 		// Send cache to client directly.