Send dns reg chunks in separate goroutines

Author: mingyech

pkg/registrars/dns-registrar/examples/client/main.go

@@ -27,16 +27,14 @@ func main() {
 	pubKey, err := hex.DecodeString(key)
 	util.Check(err)
 
-	req, err := requester.NewRequester(&requester.Config{
-		TransportMethod: requester.UDP,
-		Target:          addr.String(),
-		BaseDomain:      *baseDomain,
-		Pubkey:          pubKey,
-	})
-
-	util.Check(err)
-
-	tworeq, err := tworeqresp.NewRequester(req, 80)
+	tworeq, err := tworeqresp.NewRequester(func() (tworeqresp.Onerequester, error) {
+		return requester.NewRequester(&requester.Config{
+			TransportMethod: requester.UDP,
+			Target:          addr.String(),
+			BaseDomain:      *baseDomain,
+			Pubkey:          pubKey,
+		})
+	}, 80)
 	util.Check(err)
 
 	reader := bufio.NewReader(os.Stdin)

pkg/registrars/dns-registrar/tworeqresp/requester.go

@@ -14,19 +14,20 @@ type dialFunc = func(ctx context.Context, network, addr string) (net.Conn, error
 
 const idLen = 8
 
-type onerequester interface {
+type Onerequester interface {
 	RequestAndRecv(sendBytes []byte) ([]byte, error)
 	Close() error
 	SetDialer(dialer dialFunc) error
 }
 
 type Requester struct {
-	parent onerequester
-	mtu    uint
+	createRequester func() (Onerequester, error)
+	mtu             uint
+	dialer          dialFunc
 }
 
-func NewRequester(parent onerequester, mtu uint) (*Requester, error) {
-	return &Requester{parent: parent, mtu: mtu}, nil
+func NewRequester(createRequester func() (Onerequester, error), mtu uint) (*Requester, error) {
+	return &Requester{createRequester: createRequester, mtu: mtu}, nil
 }
 
 func (r *Requester) RequestAndRecv(sendBytes []byte) ([]byte, error) {
@@ -39,32 +40,71 @@ func (r *Requester) RequestAndRecv(sendBytes []byte) ([]byte, error) {
 
 	parts := splitIntoChunks(sendBytes, int(r.mtu))
 
-	for i, partBytes := range parts {
-		toSend := &pb.DnsPartReq{Id: id[:], PartNum: proto.Uint32(uint32(i)), TotalParts: proto.Uint32(uint32(len(parts))), Data: partBytes}
-		toSendBytes, err := proto.Marshal(toSend)
-		if err != nil {
-			return nil, fmt.Errorf("error marshal part %v: %v", i, err)
-		}
+	resCh := make(chan []byte, len(parts))
+	errCh := make(chan error, len(parts))
+	waitCh := make(chan struct{}, len(parts))
 
-		respBytes, err := r.parent.RequestAndRecv(toSendBytes)
-		if err != nil {
-			return nil, fmt.Errorf("error request part %v: %v", i, err)
-		}
+	for i, partBytes := range parts {
+		i := i
+		partBytes := partBytes
+		go func() {
+			toSend := &pb.DnsPartReq{Id: id[:], PartNum: proto.Uint32(uint32(i)), TotalParts: proto.Uint32(uint32(len(parts))), Data: partBytes}
+			toSendBytes, err := proto.Marshal(toSend)
+			if err != nil {
+				errCh <- fmt.Errorf("error marshal part %v: %v", i, err)
+				return
+			}
+
+			req, err := r.createRequester()
+			if err != nil {
+				errCh <- fmt.Errorf("error creating requester in part %v: %v", i, err)
+				return
+			}
+
+			if r.dialer != nil {
+				err = req.SetDialer(r.dialer)
+				if err != nil {
+					errCh <- fmt.Errorf("error setting dialer in part %v: %v", i, err)
+					return
+				}
+			}
+
+			respBytes, err := req.RequestAndRecv(toSendBytes)
+			if err != nil {
+				errCh <- fmt.Errorf("error request part %v: %v", i, err)
+				return
+			}
+
+			resp := &pb.DnsPartResp{}
+			err = proto.Unmarshal(respBytes, resp)
+			if err != nil {
+				errCh <- fmt.Errorf("error unmarshal response: %v", err)
+				return
+			}
+
+			if resp.GetWaiting() {
+				waitCh <- struct{}{}
+				return
+			}
+
+			resCh <- resp.GetData()
+
+		}()
+	}
 
-		resp := &pb.DnsPartResp{}
-		err = proto.Unmarshal(respBytes, resp)
-		if err != nil {
-			return nil, fmt.Errorf("error unmarshal response: %v", err)
-		}
+	errs := []error{}
 
-		if resp.GetWaiting() {
-			continue
+	for range parts {
+		select {
+		case res := <-resCh:
+			return res, nil
+		case err = <-errCh:
+			errs = append(errs, err)
+		case <-waitCh:
 		}
-
-		return resp.GetData(), nil
 	}
 
-	return nil, fmt.Errorf("no response")
+	return nil, fmt.Errorf("errors occurred: %v", errs)
 }
 
 func splitIntoChunks(data []byte, mtu int) [][]byte {
@@ -81,10 +121,11 @@ func splitIntoChunks(data []byte, mtu int) [][]byte {
 
 // Close closes the parent transport
 func (r *Requester) Close() error {
-	return r.parent.Close()
+	return nil
 }
 
 // SetDialer sets the parent dialer
 func (r *Requester) SetDialer(dialer dialFunc) error {
-	return r.parent.SetDialer(dialer)
+	r.dialer = dialer
+	return nil
 }

pkg/registrars/dns-registrar/tworeqresp/requester_test.go

@@ -1,18 +1,19 @@
 package tworeqresp
 
 import (
+	"sync/atomic"
 	"testing"
 
 	pb "github.com/refraction-networking/conjure/proto"
 	"google.golang.org/protobuf/proto"
 )
 
 type mockRequester struct {
-	calls int
+	calls atomic.Uint32
 }
 
 func (r *mockRequester) RequestAndRecv([]byte) ([]byte, error) {
-	r.calls += 1
+	r.calls.Add(1)
 	return proto.Marshal(&pb.DnsPartResp{Waiting: proto.Bool(true)})
 }
 
@@ -47,13 +48,13 @@ func TestSpliting(t *testing.T) {
 	} {
 		t.Run(testCase.name, func(t *testing.T) {
 			parent := &mockRequester{}
-			requester, err := NewRequester(parent, mtu)
+			requester, err := NewRequester(func() (Onerequester, error) { return parent, nil }, mtu)
 			if err != nil {
 				t.Fatalf("error creating requester: %v", err)
 			}
 			_, _ = requester.RequestAndRecv(testCase.data)
-			if parent.calls != testCase.chunksExpected {
-				t.Fatalf("calls: %v, expected: %v", parent.calls, testCase.chunksExpected)
+			if int(parent.calls.Load()) != testCase.chunksExpected {
+				t.Fatalf("calls: %v, expected: %v", parent.calls.Load(), testCase.chunksExpected)
 			}
 		})
 

pkg/registrars/registration/dns-registrar.go

@@ -61,16 +61,14 @@ func createRequester(config *Config) (*requester.Requester, error) {
 
 // NewDNSRegistrar creates a DNSRegistrar from config
 func NewDNSRegistrar(config *Config) (*DNSRegistrar, error) {
-	req, err := createRequester(config)
-	if err != nil {
-		return nil, fmt.Errorf("error creating requester: %v", err)
-	}
 
 	mtu := config.DNSRegistrarMTU
 	if mtu == 0 {
 		mtu = defaultMTU
 	}
-	tworeq, err := tworeqresp.NewRequester(req, mtu)
+	tworeq, err := tworeqresp.NewRequester(func() (tworeqresp.Onerequester, error) {
+		return createRequester(config)
+	}, mtu)
 	if err != nil {
 		return nil, fmt.Errorf("error adding fragmentation layer: %v", err)
 	}