Add concurrency test

Use 10 concurrent clients to send http request to remote server and
download 10M bytes content from the server.

We could test protocols other than http as well.

Fix Issue 26

Author: anfernee

pkg/agent/client/client.go

@@ -22,6 +22,7 @@ import (
 	"io"
 	"math/rand"
 	"net"
+	"sync"
 	"time"
 
 	"google.golang.org/grpc"
@@ -43,9 +44,11 @@ type dialResult struct {
 
 // grpcTunnel implements Tunnel
 type grpcTunnel struct {
-	stream      agent.ProxyService_ProxyClient
-	pendingDial map[int64]chan<- dialResult
-	conns       map[int64]*conn
+	stream          agent.ProxyService_ProxyClient
+	pendingDial     map[int64]chan<- dialResult
+	conns           map[int64]*conn
+	pendingDialLock sync.RWMutex
+	connsLock       sync.RWMutex
 }
 
 // CreateGrpcTunnel creates a Tunnel to dial to a remote server through a
@@ -90,7 +93,11 @@ func (t *grpcTunnel) serve() {
 		switch pkt.Type {
 		case agent.PacketType_DIAL_RSP:
 			resp := pkt.GetDialResponse()
-			if ch, ok := t.pendingDial[resp.Random]; !ok {
+			t.pendingDialLock.RLock()
+			ch, ok := t.pendingDial[resp.Random]
+			t.pendingDialLock.RUnlock()
+
+			if !ok {
 				klog.Warning("DialResp not recognized; dropped")
 			} else {
 				ch <- dialResult{
@@ -101,18 +108,28 @@ func (t *grpcTunnel) serve() {
 		case agent.PacketType_DATA:
 			resp := pkt.GetData()
 			// TODO: flow control
-			if conn, ok := t.conns[resp.ConnectID]; ok {
+			t.connsLock.RLock()
+			conn, ok := t.conns[resp.ConnectID]
+			t.connsLock.RUnlock()
+
+			if ok {
 				conn.readCh <- resp.Data
 			} else {
 				klog.Warningf("connection id %d not recognized", resp.ConnectID)
 			}
 		case agent.PacketType_CLOSE_RSP:
 			resp := pkt.GetCloseResponse()
-			if conn, ok := t.conns[resp.ConnectID]; ok {
+			t.connsLock.RLock()
+			conn, ok := t.conns[resp.ConnectID]
+			t.connsLock.RUnlock()
+
+			if ok {
 				close(conn.readCh)
 				conn.closeCh <- resp.Error
 				close(conn.closeCh)
+				t.connsLock.Lock()
 				delete(t.conns, resp.ConnectID)
+				t.connsLock.Unlock()
 			} else {
 				klog.Warningf("connection id %d not recognized", resp.ConnectID)
 			}
@@ -129,9 +146,13 @@ func (t *grpcTunnel) Dial(protocol, address string) (net.Conn, error) {
 
 	random := rand.Int63()
 	resCh := make(chan dialResult)
+	t.pendingDialLock.Lock()
 	t.pendingDial[random] = resCh
+	t.pendingDialLock.Unlock()
 	defer func() {
+		t.pendingDialLock.Lock()
 		delete(t.pendingDial, random)
+		t.pendingDialLock.Unlock()
 	}()
 
 	req := &agent.Packet{
@@ -163,7 +184,9 @@ func (t *grpcTunnel) Dial(protocol, address string) (net.Conn, error) {
 		c.connID = res.connid
 		c.readCh = make(chan []byte, 10)
 		c.closeCh = make(chan string)
+		t.connsLock.Lock()
 		t.conns[res.connid] = c
+		t.connsLock.Unlock()
 	case <-time.After(30 * time.Second):
 		return nil, errors.New("dial timeout")
 	}

tests/concurrent_test.go

@@ -0,0 +1,71 @@
+package tests
+
+import (
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"sigs.k8s.io/apiserver-network-proxy/pkg/agent/client"
+)
+
+func TestProxy_Concurrency(t *testing.T) {
+	length := 1 << 20
+	chunks := 10
+	server := httptest.NewServer(newSizedServer(length, chunks))
+	defer server.Close()
+
+	stopCh := make(chan struct{})
+	defer close(stopCh)
+
+	proxy, cleanup, err := runGRPCProxyServer()
+	defer cleanup()
+
+	if err := runAgent(proxy.agent, stopCh); err != nil {
+		t.Fatal(err)
+	}
+
+	// Wait for agent to register on proxy server
+	time.Sleep(time.Second)
+
+	// run test client
+	tunnel, err := client.CreateGrpcTunnel(proxy.front, grpc.WithInsecure())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var wg sync.WaitGroup
+	verify := func() {
+		defer wg.Done()
+
+		c := &http.Client{
+			Transport: &http.Transport{
+				Dial: tunnel.Dial,
+			},
+		}
+
+		r, err := c.Get(server.URL)
+		if err != nil {
+			t.Error(err)
+		}
+
+		data, err := ioutil.ReadAll(r.Body)
+		if err != nil {
+			t.Error(err)
+		}
+
+		if len(data) != length*chunks {
+			t.Errorf("expect data length %d; got %d", length*chunks, len(data))
+		}
+	}
+
+	concurrency := 10
+	wg.Add(concurrency)
+	for i := 0; i < concurrency; i++ {
+		go verify()
+	}
+	wg.Wait()
+}

tests/proxy_test.go

@@ -21,14 +21,32 @@ import (
 
 // test remote server
 type testServer struct {
+	echo   []byte
+	chunks int
+}
+
+func newEchoServer(echo string) *testServer {
+	return &testServer{
+		echo:   []byte(echo),
+		chunks: 1,
+	}
+}
+
+func newSizedServer(length, chunks int) *testServer {
+	return &testServer{
+		echo:   make([]byte, length),
+		chunks: chunks,
+	}
 }
 
 func (s *testServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	w.Write([]byte("hello"))
+	for i := 0; i < s.chunks; i++ {
+		w.Write(s.echo)
+	}
 }
 
 func TestBasicProxy_GRPC(t *testing.T) {
-	server := httptest.NewServer(&testServer{})
+	server := httptest.NewServer(newEchoServer("hello"))
 	defer server.Close()
 
 	stopCh := make(chan struct{})
@@ -74,8 +92,54 @@ func TestBasicProxy_GRPC(t *testing.T) {
 	}
 }
 
+func TestProxy_LargeResponse(t *testing.T) {
+	length := 1 << 20 // 1M
+	chunks := 10
+	server := httptest.NewServer(newSizedServer(length, chunks))
+	defer server.Close()
+
+	stopCh := make(chan struct{})
+	defer close(stopCh)
+
+	proxy, cleanup, err := runGRPCProxyServer()
+	defer cleanup()
+
+	if err := runAgent(proxy.agent, stopCh); err != nil {
+		t.Fatal(err)
+	}
+
+	// Wait for agent to register on proxy server
+	time.Sleep(time.Second)
+
+	// run test client
+	tunnel, err := client.CreateGrpcTunnel(proxy.front, grpc.WithInsecure())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	c := &http.Client{
+		Transport: &http.Transport{
+			Dial: tunnel.Dial,
+		},
+	}
+
+	r, err := c.Get(server.URL)
+	if err != nil {
+		t.Error(err)
+	}
+
+	data, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if len(data) != length*chunks {
+		t.Errorf("expect data length %d; got %d", length*chunks, len(data))
+	}
+}
+
 func TestBasicProxy_HTTPCONN(t *testing.T) {
-	server := httptest.NewServer(&testServer{})
+	server := httptest.NewServer(newEchoServer("hello"))
 	defer server.Close()
 
 	stopCh := make(chan struct{})