Merge pull request #78 from caesarxuchao/add-packet-size

Fix races

Author: k8s-ci-robot

go.sum

@@ -60,6 +60,7 @@ github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5y
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
+github.com/google/go-cmp v0.3.0 h1:crn/baboCvb5fXaQ0IJ1SGTsTVrWpDsCWC8EGETZijY=
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/gofuzz v0.0.0-20161122191042-44d81051d367/go.mod h1:HP5RmnzzSNb993RKQDq4+1A4ia9nllfqcQFTQJedwGI=
 github.com/google/gofuzz v1.0.0 h1:A8PeW59pxE9IoFRqBp37U+mSNaQoZ46F1f0f863XSXw=

pkg/agent/agentclient/client.go

@@ -74,7 +74,7 @@ func (c *connContext) cleanup() {
 	c.cleanOnce.Do(c.cleanFunc)
 }
 
-// Connect connnects to proxy server to establish a gRPC stream,
+// Connect connects to proxy server to establish a gRPC stream,
 // on which the proxied traffic is multiplexed through the stream
 // and piped to the local connection. It register itself as a
 // backend from proxy server, so proxy server will route traffic
@@ -230,7 +230,7 @@ func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64) {
 
 	for {
 		n, err := conn.Read(buf[:])
-		klog.Infof("received %d bytes from proxy server", n)
+		klog.Infof("received %d bytes from remote for connID[%d]", n, connID)
 
 		if err == io.EOF {
 			klog.Info("connection EOF")
@@ -264,8 +264,10 @@ func (a *AgentClient) proxyToRemote(conn net.Conn, connID int64) {
 		for {
 			n, err := conn.Write(d[pos:])
 			if err == nil {
+				klog.Infof("[connID: %d] write last %d data to remote", connID, n)
 				break
 			} else if n > 0 {
+				klog.Infof("[connID: %d] write %d data to remote with error: %v", connID, n, err)
 				pos += n
 			} else {
 				klog.Errorf("conn write error: %v", err)

pkg/agent/agentserver/backend_manager.go

@@ -17,19 +17,51 @@ limitations under the License.
 package agentserver
 
 import (
+	"context"
 	"math/rand"
 	"sync"
 	"time"
 
 	"k8s.io/klog"
+	client "sigs.k8s.io/apiserver-network-proxy/konnectivity-client/proto/client"
 	"sigs.k8s.io/apiserver-network-proxy/proto/agent"
 )
 
+type Backend interface {
+	Send(p *client.Packet) error
+	Context() context.Context
+}
+
+var _ Backend = &backend{}
+var _Backend = new(agent.AgentService_ConnectServer)
+
+type backend struct {
+	// TODO: this is a multi-writer single-reader pattern, it's tricky to
+	// write it using channel. Let's worry about performance later.
+	mu   sync.Mutex // mu protects conn
+	conn agent.AgentService_ConnectServer
+}
+
+func (b *backend) Send(p *client.Packet) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.conn.Send(p)
+}
+
+func (b *backend) Context() context.Context {
+	// TODO: does Context require lock protection?
+	return b.conn.Context()
+}
+
+func newBackend(conn agent.AgentService_ConnectServer) *backend {
+	return &backend{conn: conn}
+}
+
 // BackendManager is an interface to manage backend connections, i.e.,
 // connection to the proxy agents.
 type BackendManager interface {
 	// Backend returns a single backend.
-	Backend() (agent.AgentService_ConnectServer, error)
+	Backend() (Backend, error)
 	// AddBackend adds a backend.
 	AddBackend(agentID string, conn agent.AgentService_ConnectServer)
 	// RemoveBackend removes a backend.
@@ -45,7 +77,7 @@ type DefaultBackendManager struct {
 	// For a given agent, ProxyServer prefers backends[agentID][0] to send
 	// traffic, because backends[agentID][1:] are more likely to be closed
 	// by the agent to deduplicate connections to the same server.
-	backends map[string][]agent.AgentService_ConnectServer
+	backends map[string][]*backend
 	// agentID is tracked in this slice to enable randomly picking an
 	// agentID in the Backend() method. There is no reliable way to
 	// randomly pick a key from a map (in this case, the backends) in
@@ -57,7 +89,7 @@ type DefaultBackendManager struct {
 // NewDefaultBackendManager returns a DefaultBackendManager.
 func NewDefaultBackendManager() *DefaultBackendManager {
 	return &DefaultBackendManager{
-		backends: make(map[string][]agent.AgentService_ConnectServer),
+		backends: make(map[string][]*backend),
 		random:   rand.New(rand.NewSource(time.Now().UnixNano())),
 	}
 }
@@ -70,15 +102,15 @@ func (s *DefaultBackendManager) AddBackend(agentID string, conn agent.AgentServi
 	_, ok := s.backends[agentID]
 	if ok {
 		for _, v := range s.backends[agentID] {
-			if v == conn {
+			if v.conn == conn {
 				klog.Warningf("this should not happen. Adding existing connection %v for agentID %s", conn, agentID)
 				return
 			}
 		}
-		s.backends[agentID] = append(s.backends[agentID], conn)
+		s.backends[agentID] = append(s.backends[agentID], newBackend(conn))
 		return
 	}
-	s.backends[agentID] = []agent.AgentService_ConnectServer{conn}
+	s.backends[agentID] = []*backend{newBackend(conn)}
 	s.agentIDs = append(s.agentIDs, agentID)
 }
 
@@ -94,7 +126,7 @@ func (s *DefaultBackendManager) RemoveBackend(agentID string, conn agent.AgentSe
 	}
 	var found bool
 	for i, c := range backends {
-		if c == conn {
+		if c.conn == conn {
 			s.backends[agentID] = append(s.backends[agentID][:i], s.backends[agentID][i+1:]...)
 			if i == 0 && len(s.backends) != 0 {
 				klog.Warningf("this should not happen. Removed connection %v that is not the first connection, remaining connections are %v", conn, s.backends[agentID])
@@ -126,13 +158,14 @@ func (e *ErrNotFound) Error() string {
 }
 
 // Backend returns a random backend.
-func (s *DefaultBackendManager) Backend() (agent.AgentService_ConnectServer, error) {
+func (s *DefaultBackendManager) Backend() (Backend, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 	if len(s.backends) == 0 {
 		return nil, &ErrNotFound{}
 	}
 	agentID := s.agentIDs[s.random.Intn(len(s.agentIDs))]
+	klog.Infof("pick agentID=%s as backend", agentID)
 	// always return the first connection to an agent, because the agent
 	// will close later connections if there are multiple.
 	return s.backends[agentID][0], nil

pkg/agent/agentserver/backend_manager_test.go

@@ -38,7 +38,7 @@ func TestAddRemoveBackends(t *testing.T) {
 
 	p.AddBackend("agent1", conn1)
 	p.RemoveBackend("agent1", conn1)
-	expectedBackends := make(map[string][]agent.AgentService_ConnectServer)
+	expectedBackends := make(map[string][]*backend)
 	expectedAgentIDs := []string{}
 	if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) {
 		t.Errorf("expected %v, got %v", e, a)
@@ -60,9 +60,9 @@ func TestAddRemoveBackends(t *testing.T) {
 	p.RemoveBackend("agent1", conn1)
 	// This is invalid. agent1 doesn't have conn3. This should be a no-op.
 	p.RemoveBackend("agent1", conn3)
-	expectedBackends = map[string][]agent.AgentService_ConnectServer{
-		"agent1": []agent.AgentService_ConnectServer{conn12},
-		"agent3": []agent.AgentService_ConnectServer{conn3},
+	expectedBackends = map[string][]*backend{
+		"agent1": []*backend{newBackend(conn12)},
+		"agent3": []*backend{newBackend(conn3)},
 	}
 	expectedAgentIDs = []string{"agent1", "agent3"}
 	if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) {

pkg/agent/agentserver/server.go

@@ -41,6 +41,7 @@ type ProxyClientConnection struct {
 	HTTP      net.Conn
 	connected chan struct{}
 	connectID int64
+	agentID   string
 }
 
 func (c *ProxyClientConnection) send(pkt *client.Packet) error {
@@ -63,6 +64,36 @@ func (c *ProxyClientConnection) send(pkt *client.Packet) error {
 	}
 }
 
+func NewPendingDialManager() *PendingDialManager {
+	return &PendingDialManager{
+		pendingDial: make(map[int64]*ProxyClientConnection),
+	}
+}
+
+type PendingDialManager struct {
+	mu          sync.RWMutex
+	pendingDial map[int64]*ProxyClientConnection
+}
+
+func (pm *PendingDialManager) Add(random int64, clientConn *ProxyClientConnection) {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+	pm.pendingDial[random] = clientConn
+}
+
+func (pm *PendingDialManager) Get(random int64) (*ProxyClientConnection, bool) {
+	pm.mu.RLock()
+	defer pm.mu.RUnlock()
+	clientConn, ok := pm.pendingDial[random]
+	return clientConn, ok
+}
+
+func (pm *PendingDialManager) Remove(random int64) {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+	delete(pm.pendingDial, random)
+}
+
 // ProxyServer
 type ProxyServer struct {
 	// BackendManager manages the backends.
@@ -73,7 +104,7 @@ type ProxyServer struct {
 	// conn = Frontend[agentID][connID]
 	frontends map[string]map[int64]*ProxyClientConnection
 
-	PendingDial map[int64]*ProxyClientConnection
+	PendingDial *PendingDialManager
 
 	serverID    string // unique ID of this server
 	serverCount int    // Number of proxy server instances, should be 1 unless it is a HA server.
@@ -143,7 +174,7 @@ func (s *ProxyServer) getFrontend(agentID string, connID int64) (*ProxyClientCon
 func NewProxyServer(serverID string, serverCount int, agentAuthenticationOptions *AgentTokenAuthenticationOptions) *ProxyServer {
 	p := &ProxyServer{
 		frontends:                  make(map[string](map[int64]*ProxyClientConnection)),
-		PendingDial:                make(map[int64]*ProxyClientConnection),
+		PendingDial:                NewPendingDialManager(),
 		serverID:                   serverID,
 		serverCount:                serverCount,
 		BackendManager:             NewDefaultBackendManager(),
@@ -192,7 +223,7 @@ func (s *ProxyServer) serveRecvFrontend(stream client.ProxyService_ProxyServer,
 	var firstConnID int64
 	// The first packet should be a DIAL_REQ, we will randomly get a
 	// backend from the BackendManger then.
-	var backend agent.AgentService_ConnectServer
+	var backend Backend
 	var err error
 
 	for pkt := range recvCh {
@@ -211,11 +242,13 @@ func (s *ProxyServer) serveRecvFrontend(stream client.ProxyService_ProxyServer,
 			if err := backend.Send(pkt); err != nil {
 				klog.Warningf(">>> DIAL_REQ to Backend failed: %v", err)
 			}
-			s.PendingDial[pkt.GetDialRequest().Random] = &ProxyClientConnection{
-				Mode:      "grpc",
-				Grpc:      stream,
-				connected: make(chan struct{}),
-			}
+			s.PendingDial.Add(
+				pkt.GetDialRequest().Random,
+				&ProxyClientConnection{
+					Mode:      "grpc",
+					Grpc:      stream,
+					connected: make(chan struct{}),
+				})
 			klog.Info(">>> DIAL_REQ sent to backend") // got this. but backend didn't receive anything.
 
 		case client.PacketType_CLOSE_REQ:
@@ -233,7 +266,8 @@ func (s *ProxyServer) serveRecvFrontend(stream client.ProxyService_ProxyServer,
 
 		case client.PacketType_DATA:
 			connID := pkt.GetData().ConnectID
-			klog.Infof(">>> Received DATA(id=%d)", connID)
+			data := pkt.GetData().Data
+			klog.Infof(">>> Received %d bytes of DATA(id=%d)", len(data), connID)
 			if firstConnID == 0 {
 				firstConnID = connID
 			} else if firstConnID != connID {
@@ -424,25 +458,26 @@ func (s *ProxyServer) serveRecvBackend(stream agent.AgentService_ConnectServer,
 		switch pkt.Type {
 		case client.PacketType_DIAL_RSP:
 			resp := pkt.GetDialResponse()
-			klog.Infof("<<< Received DIAL_RSP(rand=%d, id=%d)", resp.Random, resp.ConnectID)
+			klog.Infof("<<< Received DIAL_RSP(rand=%d), agentID %s, connID %d)", resp.Random, agentID, resp.ConnectID)
 
-			if client, ok := s.PendingDial[resp.Random]; !ok {
+			if client, ok := s.PendingDial.Get(resp.Random); !ok {
 				klog.Warning("<<< DialResp not recognized; dropped")
 			} else {
 				err := client.send(pkt)
-				delete(s.PendingDial, resp.Random)
+				s.PendingDial.Remove(resp.Random)
 				if err != nil {
 					klog.Warningf("<<< DIAL_RSP send to client stream error: %v", err)
 				} else {
 					client.connectID = resp.ConnectID
+					client.agentID = agentID
 					s.addFrontend(agentID, resp.ConnectID, client)
 					close(client.connected)
 				}
 			}
 
 		case client.PacketType_DATA:
 			resp := pkt.GetData()
-			klog.Infof("<<< Received DATA(id=%d)", resp.ConnectID)
+			klog.Infof("<<< Received %d bytes of DATA from agentID %s, connID %d", len(resp.Data), agentID, resp.ConnectID)
 			client, err := s.getFrontend(agentID, resp.ConnectID)
 			if err != nil {
 				klog.Warning(err)
@@ -469,7 +504,7 @@ func (s *ProxyServer) serveRecvBackend(stream agent.AgentService_ConnectServer,
 				klog.Infof("<<< CLOSE_RSP sent to frontend")
 			}
 			s.removeFrontend(agentID, resp.ConnectID)
-			klog.Infof("<<< Close streaming (id=%d)", resp.ConnectID)
+			klog.Infof("<<< Close streaming (agentID=%s, connId=%d)", agentID, resp.ConnectID)
 
 		default:
 			klog.Warningf("<<< Unrecognized packet %+v", pkt)

pkg/agent/agentserver/tunnel.go

@@ -73,7 +73,7 @@ func (t *Tunnel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		HTTP:      conn,
 		connected: connected,
 	}
-	t.Server.PendingDial[random] = connection
+	t.Server.PendingDial.Add(random, connection)
 	backend, err := t.Server.BackendManager.Backend()
 	if err != nil {
 		http.Error(w, fmt.Sprintf("currently no tunnels available: %v", err), http.StatusInternalServerError)
@@ -102,8 +102,13 @@ func (t *Tunnel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	klog.Infof("Starting proxy to %q", r.Host)
 	pkt := make([]byte, 1<<12)
 
+	connID := connection.connectID
+	agentID := connection.agentID
+	var acc int
+
 	for {
 		n, err := bufrw.Read(pkt[:])
+		acc += n
 		if err == io.EOF {
 			klog.Warningf("EOF from %v", r.Host)
 			break
@@ -117,7 +122,7 @@ func (t *Tunnel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			Type: client.PacketType_DATA,
 			Payload: &client.Packet_Data{
 				Data: &client.Data{
-					ConnectID: connection.connectID,
+					ConnectID: connID,
 					Data:      pkt[:n],
 				},
 			},
@@ -127,8 +132,8 @@ func (t *Tunnel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			klog.Errorf("error sending packet %v", err)
 			continue
 		}
-		klog.Infof("Forwarding on tunnel, packet type: %s", packet.Type)
+		klog.Infof("Forwarding %d (total %d) bytes of DATA on tunnel for agentID %s, connID %d", n, acc, connection.agentID, connection.connectID)
 	}
 
-	klog.Infof("Stopping transfer to %q", r.Host)
+	klog.Infof("Stopping transfer to %q, agentID %s, connID %d", r.Host, agentID, connID)
 }

tests/concurrent_client_request_test.go

@@ -78,7 +78,7 @@ func (s *singleTimeManager) RemoveBackend(agentID string, conn agent.AgentServic
 	delete(s.backends, agentID)
 }
 
-func (s *singleTimeManager) Backend() (agent.AgentService_ConnectServer, error) {
+func (s *singleTimeManager) Backend() (agentserver.Backend, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	for k, v := range s.backends {