client: Add massively concurrent test, do not deadlock

Add a massively concurrent unit test.

Remove a chance for a deadlock due to full channel while the recipient is
blocked on a lock by only ever waking up one waiter when a matching
response has been received and unanswered requests remain.

Signed-off-by: Jarno Rajahalme <jarno@isovalent.com>

go.mod

@@ -1,6 +1,6 @@
 module github.com/cilium/dns
 
-go 1.14
+go 1.18
 
 require (
 	golang.org/x/net v0.2.0

shared_client.go

@@ -49,7 +49,6 @@ func (s *SharedClients) GetSharedClient(key string, conf *Client, serverAddrStr
 	return client, func() {
 		s.Lock()
 		defer s.Unlock()
-
 		client.refcount--
 		if client.refcount == 0 {
 			// Make client unreachable and close it's connection.
@@ -133,6 +132,11 @@ func (c *SharedClient) ExchangeSharedContext(ctx context.Context, m *Msg) (r *Ms
 		}
 	}
 
+	// Create channel for the response with buffer of one, so that write to it
+	// does not block if we happen to do it ourselves.
+	respCh := make(chan Response, 1)
+	c.reqs[m.Id] = respCh
+
 	// Send while holding the client lock, as Client is not made to be usable from
 	// concurrent goroutines.
 	start := time.Now()
@@ -141,11 +145,6 @@ func (c *SharedClient) ExchangeSharedContext(ctx context.Context, m *Msg) (r *Ms
 		return nil, 0, err
 	}
 
-	// Create channel for the response with buffer of one, so that write to it
-	// does not block if we happen to do it ourselves.
-	ch := make(chan Response, 1)
-	c.reqs[m.Id] = ch
-
 	// Wait for the response
 	var resp Response
 	for {
@@ -185,22 +184,30 @@ func (c *SharedClient) ExchangeSharedContext(ctx context.Context, m *Msg) (r *Ms
 			// Releasing the reader lock before sending errors on waiter's channels
 			// so that when they get them, one of them can take the reader lock.
 			c.readerLock.Unlock()
-			for id, ch := range c.reqs {
-				// Another reader will pick up if any errNoReader errors are sent.
-				// Only delete the pending request in other error cases.
-				if !errors.Is(err, errNoReader) {
+			if errors.Is(err, errNoReader) {
+				// Can only wake one waiting requester to do the reading as the
+				// channel buffer length is one, otherwise the channel could get
+				// full while the request is still waiting for a lock.
+				for _, ch := range c.reqs {
+					ch <- Response{err: err}
+					break
+				}
+			} else {
+				// Other errors are sent to all recipients
+				for id, ch := range c.reqs {
 					delete(c.reqs, id)
+					ch <- Response{err: err}
 				}
-				ch <- Response{err: err}
 			}
 		}
 		// Get the response of error from the current reader.
 		// Unlock for the blocking duration to allow concurrent writes
 		// on the client's connection.
 		c.Unlock()
-		resp = <-ch
+		resp = <-respCh
 		c.Lock()
 		if !errors.Is(resp.err, errNoReader) {
+			// error other than errNoReader received
 			break
 		}
 		// Trying again

shared_client_test.go

@@ -6,9 +6,11 @@ package dns
 import (
 	"context"
 	"crypto/tls"
+	"fmt"
 	"net"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 )
@@ -58,6 +60,7 @@ func TestSharedClientSync(t *testing.T) {
 	if c2 != c {
 		t.Fatal("client not really shared")
 	}
+	m.Id = uint16(42)
 	r, _, err = c2.ExchangeShared(m)
 	if err != nil {
 		t.Errorf("failed to exchange: %v", err)
@@ -67,6 +70,60 @@ func TestSharedClientSync(t *testing.T) {
 	}
 }
 
+func TestSharedClientConcurrentSync(t *testing.T) {
+	HandleFunc("miek.nl.", HelloServer)
+	defer HandleRemove("miek.nl.")
+
+	s, addrstr, _, err := RunLocalUDPServer(":0")
+	if err != nil {
+		t.Fatalf("unable to run test server: %v", err)
+	}
+	defer s.Shutdown()
+
+	conf := &Client{
+		Timeout: time.Second,
+	}
+
+	errors := sync.Map{}
+	var wg sync.WaitGroup
+	f1 := func(id uint16) {
+		defer wg.Done()
+
+		m := new(Msg)
+		m.SetQuestion("miek.nl.", TypeSOA)
+		m.Id = id
+
+		c, closer := clients.GetSharedClient("concurrent-client", conf, addrstr)
+		defer closer()
+		r, _, err := c.ExchangeShared(m)
+		if err != nil {
+			errors.Store(id, fmt.Errorf("failed to exchange: %v", err))
+		} else {
+			if r == nil {
+				errors.Store(id, fmt.Errorf("response is nil"))
+			} else {
+				if r.Id != id {
+					errors.Store(id, fmt.Errorf("incorrect id (%d != %d)", r.Id, id))
+				}
+				if r.Rcode != RcodeSuccess {
+					errors.Store(id, fmt.Errorf("failed to get an valid answer\n%v", r))
+				}
+			}
+		}
+	}
+
+	for id := uint16(1); id <= 100; id++ {
+		wg.Add(1)
+		go f1(id)
+	}
+	wg.Wait()
+
+	errors.Range(func(key, value any) bool {
+		t.Fatalf("Id: %v, error: %v", key, value)
+		return false
+	})
+}
+
 func TestSharedClientLocalAddress(t *testing.T) {
 	HandleFunc("miek.nl.", HelloServerEchoAddrPort)
 	defer HandleRemove("miek.nl.")
@@ -90,7 +147,10 @@ func TestSharedClientLocalAddress(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to exchange: %v", err)
 	}
-	if r != nil && r.Rcode != RcodeSuccess {
+	if r == nil {
+		t.Fatalf("No response")
+	}
+	if r.Rcode != RcodeSuccess {
 		t.Errorf("failed to get an valid answer\n%v", r)
 	}
 	if len(r.Extra) != 1 {
@@ -210,7 +270,10 @@ func TestSharedClientSyncBadThenGoodID(t *testing.T) {
 
 	r, _, err := c.ExchangeShared(m)
 	if err != nil {
-		t.Errorf("failed to exchange: %v", err)
+		t.Fatalf("failed to exchange: %v", err)
+	}
+	if r == nil {
+		t.Fatalf("No response")
 	}
 	if r.Id != m.Id {
 		t.Errorf("failed to get response with expected Id")
@@ -435,9 +498,12 @@ func TestSharedConcurrentExchanges(t *testing.T) {
 
 		r := make([]*Msg, 2)
 		for i := range r {
-			r[i], _, _ = c.ExchangeShared(m.Copy())
+			r[i], _, err = c.ExchangeShared(m.Copy())
+			if err != nil {
+				t.Fatalf("failed to exchange: %s", err)
+			}
 			if r[i] == nil {
-				t.Errorf("response %d is nil", i)
+				t.Fatalf("response %d is nil", i)
 			}
 		}