net: avoid relying on singleflight.Group.DoChan to detect hook called

So next CLs can revert changes to "internal/singleflight" in CL #82795,
then replace it with "golang.org/x/sync/singleflight" instead.

For #31697

Change-Id: I873ce30d7e051539aa6dc5d4f21e558869a6d132
Reviewed-on: https://go-review.googlesource.com/c/go/+/423654
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Bryan Mills <bcmills@google.com>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Cuong Manh Le <cuong.manhle.vn@gmail.com>

diff --git a/src/net/lookup.go b/src/net/lookup.go
index 7f3d20126c902c79e863f06d0f2e4fd49d8e38ee..3cc53f1db6857b0401c30ac40b8fc33709af3fc4 100644 (file)
--- a/src/net/lookup.go
+++ b/src/net/lookup.go
@@ -316,14 +316,15 @@ func (r *Resolver) lookupIPAddr(ctx context.Context, network, host string) ([]IP
 
        lookupKey := network + "\000" + host
        dnsWaitGroup.Add(1)
-       ch, called := r.getLookupGroup().DoChan(lookupKey, func() (any, error) {
-               defer dnsWaitGroup.Done()
+       ch, _ := r.getLookupGroup().DoChan(lookupKey, func() (any, error) {
                return testHookLookupIP(lookupGroupCtx, resolverFunc, network, host)
        })
-       if !called {
+
+       dnsWaitGroupDone := func(ch <-chan singleflight.Result, cancelFn context.CancelFunc) {
+               <-ch
                dnsWaitGroup.Done()
+               cancelFn()
        }
-
        select {
        case <-ctx.Done():
                // Our context was canceled. If we are the only
@@ -335,11 +336,9 @@ func (r *Resolver) lookupIPAddr(ctx context.Context, network, host string) ([]IP
                // See issues 8602, 20703, 22724.
                if r.getLookupGroup().ForgetUnshared(lookupKey) {
                        lookupGroupCancel()
+                       go dnsWaitGroupDone(ch, func() {})
                } else {
-                       go func() {
-                               <-ch
-                               lookupGroupCancel()
-                       }()
+                       go dnsWaitGroupDone(ch, lookupGroupCancel)
                }
                ctxErr := ctx.Err()
                err := &DNSError{
@@ -352,6 +351,7 @@ func (r *Resolver) lookupIPAddr(ctx context.Context, network, host string) ([]IP
                }
                return nil, err
        case r := <-ch:
+               dnsWaitGroup.Done()
                lookupGroupCancel()
                err := r.Err
                if err != nil {