Implement Close

- When agent receives CLOSE_REQ, it closes the corresponding connection
identified by ConnID
- When agent's connection failed, it sends CLOSE_RESP to server to
notify it.

Author: anfernee

pkg/agent/agentclient/client.go

@@ -29,19 +29,28 @@ import (
 )
 
 type AgentClient struct {
-	nextConnID int64
-	conns      map[int64]net.Conn
-	dataChs    map[int64]chan []byte
-	address    string
+	nextConnID  int64
+	connContext map[int64]*connContext
+	address     string
 
 	stream agent.AgentService_ConnectClient
 }
 
+type connContext struct {
+	conn      net.Conn
+	cleanFunc func()
+	dataCh    chan []byte
+	cleanOnce sync.Once
+}
+
+func (c *connContext) cleanup() {
+	c.cleanOnce.Do(c.cleanFunc)
+}
+
 func NewAgentClient(address string) *AgentClient {
 	a := &AgentClient{
-		conns:   make(map[int64]net.Conn),
-		address: address,
-		dataChs: make(map[int64]chan []byte),
+		connContext: make(map[int64]*connContext),
+		address:     address,
 	}
 
 	return a
@@ -107,17 +116,29 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 			}
 
 			connID := atomic.AddInt64(&a.nextConnID, 1)
-			a.conns[connID] = conn
-			a.dataChs[connID] = make(chan []byte, 5)
-			var once sync.Once
-
-			cleanup := func() {
-				once.Do(func() {
-					conn.Close()
-					close(a.dataChs[connID])
-					delete(a.conns, connID)
-					delete(a.dataChs, connID)
-				})
+			dataCh := make(chan []byte, 5)
+			a.connContext[connID] = &connContext{
+				conn:   conn,
+				dataCh: dataCh,
+				cleanFunc: func() {
+					resp := &agent.Packet{
+						Type:    agent.PacketType_CLOSE_RSP,
+						Payload: &agent.Packet_CloseResponse{CloseResponse: &agent.CloseResponse{}},
+					}
+					resp.GetCloseResponse().ConnectID = connID
+
+					err := conn.Close()
+					if err != nil {
+						resp.GetCloseResponse().Error = err.Error()
+					}
+
+					if err := a.stream.Send(resp); err != nil {
+						klog.Warningf("close response send error: %v", err)
+					}
+
+					close(dataCh)
+					delete(a.connContext, connID)
+				},
 			}
 
 			resp.GetDialResponse().ConnectID = connID
@@ -126,16 +147,37 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 				continue
 			}
 
-			go a.remoteToProxy(conn, connID, cleanup)
-			go a.proxyToRemote(conn, a.dataChs[connID], cleanup)
+			go a.remoteToProxy(conn, connID)
+			go a.proxyToRemote(conn, connID)
 
 		case agent.PacketType_DATA:
-			klog.Info("received DATA")
 			data := pkt.GetData()
+			klog.Infof("received DATA(id=%d)", data.ConnectID)
 			klog.Infof("[tracing] %v", data)
 
-			if dataCh, ok := a.dataChs[data.ConnectID]; ok {
-				dataCh <- data.Data
+			if ctx, ok := a.connContext[data.ConnectID]; ok {
+				ctx.dataCh <- data.Data
+			}
+
+		case agent.PacketType_CLOSE_REQ:
+			closeReq := pkt.GetCloseRequest()
+			connID := closeReq.ConnectID
+
+			klog.Infof("received CLOSE_REQ(id=%d)", connID)
+
+			if ctx, ok := a.connContext[connID]; ok {
+				ctx.cleanup()
+			} else {
+				resp := &agent.Packet{
+					Type:    agent.PacketType_CLOSE_RSP,
+					Payload: &agent.Packet_CloseResponse{CloseResponse: &agent.CloseResponse{}},
+				}
+				resp.GetCloseResponse().ConnectID = connID
+				resp.GetCloseResponse().Error = "Unknown connectID"
+				if err := a.stream.Send(resp); err != nil {
+					klog.Warningf("close response send error: %v", err)
+					continue
+				}
 			}
 
 		default:
@@ -144,8 +186,13 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 	}
 }
 
-func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64, cleanup func()) {
-	defer cleanup()
+func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64) {
+	ctx := a.connContext[connID]
+	if ctx == nil {
+		return
+	}
+
+	defer ctx.cleanup()
 
 	var buf [1 << 12]byte
 	resp := &agent.Packet{
@@ -154,17 +201,13 @@ func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64, cleanup func())
 
 	for {
 		n, err := conn.Read(buf[:])
+
 		if err == io.EOF {
+			klog.Info("connection EOF")
 			return
 		} else if err != nil {
-			klog.Errorf("connection read error: %v", err)
-			resp.Payload = &agent.Packet_Data{Data: &agent.Data{
-				Error:     err.Error(),
-				ConnectID: connID,
-			}}
-			if err := a.stream.Send(resp); err != nil {
-				klog.Warningf("stream send error: %v", err)
-			}
+			klog.Warningf("connection read error: %v", err)
+			return
 		} else {
 			resp.Payload = &agent.Packet_Data{Data: &agent.Data{
 				Data:      buf[:n],
@@ -177,10 +220,15 @@ func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64, cleanup func())
 	}
 }
 
-func (a *AgentClient) proxyToRemote(conn net.Conn, dataCh <-chan []byte, cleanup func()) {
-	defer cleanup()
+func (a *AgentClient) proxyToRemote(conn net.Conn, connID int64) {
+	ctx := a.connContext[connID]
+	if ctx == nil {
+		return
+	}
+
+	defer ctx.cleanup()
 
-	for d := range dataCh {
+	for d := range ctx.dataCh {
 		pos := 0
 		for {
 			n, err := conn.Write(d[pos:])