feat(worker): Report bytes up/down on Worker connection closure

Introduces data about bytes read and written to the connection close
payload. No changes are required in the Controller because the RPC
messages already contained the required fields. As of this commit,
bytes_up and bytes_down columns in the session_connection table are now
populated when a connection closes.

Refs: #2509

internal/daemon/worker/countingconn.go

@@ -5,6 +5,10 @@ import (
 	"sync"
 )
 
+// countingConn is a `net.Conn` implementation that records the bytes that go
+// across Read() and Write(). All other `net.Conn` function calls are a
+// pass-through to the underlying `net.Conn`, meaning it's also safe to call
+// those functions directly on the underlying object, if you have access to it.
 type countingConn struct {
 	net.Conn
 
@@ -29,7 +33,8 @@ func (c *countingConn) BytesWritten() uint64 {
 	return c.bytesWritten
 }
 
-// Read wraps the embedded conn's Read and counts the number of bytes read.
+// Read wraps the embedded conn's Read() and counts the number of bytes read
+// (the number of bytes the client sent to us).
 func (c *countingConn) Read(in []byte) (int, error) {
 	n, err := c.Conn.Read(in)
 	c.mu.Lock()
@@ -38,7 +43,8 @@ func (c *countingConn) Read(in []byte) (int, error) {
 	return n, err
 }
 
-// Write wraps the embedded conn's Write and counts the number of bytes read.
+// Write wraps the embedded conn's Write() and counts the number of bytes
+// written (the number of bytes we sent to the client).
 func (c *countingConn) Write(in []byte) (int, error) {
 	n, err := c.Conn.Write(in)
 	c.mu.Lock()

internal/daemon/worker/handler.go

@@ -219,9 +219,28 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig, sessionMa
 		}
 		event.WriteSysEvent(ctx, op, "connection successfully authorized", "session_id", sessionId, "connection_id", ci.Id)
 
+		// Wrapping the client websocket with a `net.Conn` implementation that
+		// records the bytes that go across Read() and Write().
 		cc := &countingConn{Conn: websocket.NetConn(connCtx, conn, websocket.MessageBinary)}
+		err = sess.ApplyConnectionCounterCallbacks(ci.Id, cc.BytesRead, cc.BytesWritten)
+		if err != nil {
+			event.WriteError(ctx, op, err, event.WithInfoMsg("unable to set counter callbacks for session connection"))
+			err = conn.Close(websocket.StatusInternalError, "unable to set counter callbacks for session connection")
+			if err != nil {
+				event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection"))
+			}
+			return
+		}
+
 		defer func() {
-			if sessionManager.RequestCloseConnections(ctx, map[string]string{ci.Id: sess.GetId()}) {
+			ccd := map[string]*session.ConnectionCloseData{
+				ci.Id: {
+					SessionId: sess.GetId(),
+					BytesUp:   cc.BytesRead(),
+					BytesDown: cc.BytesWritten(),
+				},
+			}
+			if sessionManager.RequestCloseConnections(ctx, ccd) {
 				event.WriteSysEvent(ctx, op, "connection closed", "session_id", sessionId, "connection_id", ci.Id)
 			}
 		}()

internal/daemon/worker/session/manager.go

@@ -35,15 +35,15 @@ type Manager interface {
 	// no error is returned.
 	DeleteLocalSession([]string)
 
-	// RequestCloseConnections sends connection close requests to the controller,
-	// and sets close times within the worker. It should be called during the worker
-	// status loop and on connection exit on the proxy.
+	// RequestCloseConnections sends connection close requests to the
+	// controller, and sets close times within the worker. It should be called
+	// during the worker status loop and on connection exit on the proxy.
 	//
 	// The boolean indicates whether the function was successful, e.g. had any
 	// errors. Individual events will be sent for the errors if there are any.
 	//
-	// closeInfo is a map of connection ids mapped to their individual session ids.
-	RequestCloseConnections(context.Context, map[string]string) bool
+	// closeInfo is a map of connection ids mapped to connection metadata.
+	RequestCloseConnections(context.Context, map[string]*ConnectionCloseData) bool
 }
 
 type manager struct {
@@ -123,7 +123,7 @@ func (m *manager) DeleteLocalSession(sessIds []string) {
 	}
 }
 
-func (m *manager) RequestCloseConnections(ctx context.Context, closeInfo map[string]string) bool {
+func (m *manager) RequestCloseConnections(ctx context.Context, closeInfo map[string]*ConnectionCloseData) bool {
 	return closeConnections(ctx, m.controllerSessionConn, m, closeInfo)
 }
 

internal/daemon/worker/session/manager_test.go

@@ -75,7 +75,7 @@ func TestManager_RequestCloseConnections(t *testing.T) {
 	manager, err := NewManager(mockSessionClient)
 	require.NoError(t, err)
 	assert.False(t, manager.RequestCloseConnections(ctx, nil))
-	assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{}))
+	assert.False(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{}))
 
 	mockSessionClient.LookupSessionFn = func(_ context.Context, req *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) {
 		return &pbs.LookupSessionResponse{
@@ -98,7 +98,7 @@ func TestManager_RequestCloseConnections(t *testing.T) {
 	mockSessionClient.CloseConnectionFn = func(context.Context, *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) {
 		return nil, errors.New("error")
 	}
-	assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{"connection id": session1.GetId()}))
+	assert.False(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{"connection id": {SessionId: session1.GetId()}}))
 	mockSessionClient.CloseConnectionFn = func(_ context.Context, req *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) {
 		var data []*pbs.CloseConnectionResponseData
 		for _, r := range req.GetCloseRequestData() {
@@ -110,8 +110,8 @@ func TestManager_RequestCloseConnections(t *testing.T) {
 		return &pbs.CloseConnectionResponse{CloseResponseData: data}, nil
 	}
 	// There is no connection information yet for this connection.
-	assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{
-		"random_connection_id": session1.GetId(),
+	assert.False(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{
+		"random_connection_id": {SessionId: session1.GetId()},
 	}))
 
 	// Load the connection info into the local storage
@@ -126,17 +126,17 @@ func TestManager_RequestCloseConnections(t *testing.T) {
 	c1, _, err := session1.RequestAuthorizeConnection(ctx, "worker id", cancelFn)
 	require.NoError(t, err)
 
-	assert.True(t, manager.RequestCloseConnections(ctx, map[string]string{
-		c1.Id: session1.GetId(),
+	assert.True(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{
+		c1.Id: {SessionId: session1.GetId()},
 	}))
 
 	c2, _, err := session2.RequestAuthorizeConnection(ctx, "worker id", cancelFn)
 	require.NoError(t, err)
 	c3, _, err := session3.RequestAuthorizeConnection(ctx, "worker id", cancelFn)
 	require.NoError(t, err)
-	assert.True(t, manager.RequestCloseConnections(ctx, map[string]string{
-		c2.Id: session2.GetId(),
-		c3.Id: session3.GetId(),
+	assert.True(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{
+		c2.Id: {SessionId: session2.GetId()},
+		c3.Id: {SessionId: session3.GetId()},
 	}))
 }
 

internal/daemon/worker/session/session.go

@@ -30,11 +30,24 @@ type ConnInfo struct {
 	// closes the proxy connection.
 	connCtxCancelFunc context.CancelFunc
 
+	// The number of bytes uploaded from the client.
+	BytesUp func() uint64
+	// The number of bytes downloaded to the client.
+	BytesDown func() uint64
+
 	// The time the controller has successfully reported that this connection is
 	// closed.
 	CloseTime time.Time
 }
 
+// ConnectionCloseData encapsulates the data we need to send via CloseConnection
+// RPC to the controller.
+type ConnectionCloseData struct {
+	SessionId string
+	BytesUp   uint64
+	BytesDown uint64
+}
+
 // Session is the local representation of a session.  After initial loading
 // the only values that will change will be the status (readable from
 // GetStatus()) and the Connections (GetLocalConnections()).
@@ -86,6 +99,12 @@ type Session interface {
 	// call to RequestActivate, subsequent calls will fail.
 	RequestActivate(ctx context.Context, tofu string) error
 
+	// ApplyConnectionCounterCallbacks sets a connection's bytes up and bytes
+	// down callbacks to the provided functions. Both functions must be safe for
+	// concurrent use. If there is no connection with the provided id, an error
+	// is returned.
+	ApplyConnectionCounterCallbacks(connId string, bytesUp func() uint64, bytesDown func() uint64) error
+
 	// RequestAuthorizeConnection sends an AuthorizeConnection request to
 	// the controller.
 	// It is called by the worker handler after a connection has been received by
@@ -179,6 +198,8 @@ func (s *sess) GetLocalConnections() map[string]ConnInfo {
 			Id:        v.Id,
 			Status:    v.Status,
 			CloseTime: v.CloseTime,
+			BytesUp:   v.BytesUp,
+			BytesDown: v.BytesDown,
 		}
 	}
 	return res
@@ -283,11 +304,19 @@ func (s *sess) RequestAuthorizeConnection(ctx context.Context, workerId string,
 	case workerId == "":
 		return ConnInfo{}, 0, errors.New("worker id is empty")
 	}
+
 	ci, connsLeft, err := authorizeConnection(ctx, s.client, workerId, s.GetId())
 	if err != nil {
 		return ConnInfo{}, connsLeft, err
 	}
 	ci.connCtxCancelFunc = connCancel
+
+	// Install safe callbacks before connection has been established. These
+	// should be replaced when `ApplyConnectionCounterCallbacks` gets called on
+	// the `sess` object.
+	ci.BytesUp = func() uint64 { return 0 }
+	ci.BytesDown = func() uint64 { return 0 }
+
 	s.lock.Lock()
 	defer s.lock.Unlock()
 	s.connInfoMap[ci.Id] = ci
@@ -295,6 +324,8 @@ func (s *sess) RequestAuthorizeConnection(ctx context.Context, workerId string,
 		Id:        ci.Id,
 		Status:    ci.Status,
 		CloseTime: ci.CloseTime,
+		BytesUp:   ci.BytesUp,
+		BytesDown: ci.BytesDown,
 	}, connsLeft, err
 }
 
@@ -355,6 +386,19 @@ func (s *sess) CancelAllLocalConnections() []string {
 	return closedIds
 }
 
+func (s *sess) ApplyConnectionCounterCallbacks(connId string, bytesUp func() uint64, bytesDown func() uint64) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	ci, ok := s.connInfoMap[connId]
+	if !ok {
+		return fmt.Errorf("failed to find connection info for connection id %q", connId)
+	}
+	ci.BytesUp = bytesUp
+	ci.BytesDown = bytesDown
+	return nil
+}
+
 // activate is a helper worker function that sends session activation request to the
 // controller.
 func activate(ctx context.Context, sessClient pbs.SessionServiceClient, sessionId, tofuToken string, version uint32) (pbs.SESSIONSTATUS, error) {
@@ -437,8 +481,8 @@ func closeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, r
 // The boolean indicates whether the function was successful, e.g. had any
 // errors. Individual events will be sent for the errors if there are any.
 //
-// closeInfo is a map of connections mapped to their individual session.
-func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient, sManager Manager, closeInfo map[string]string) bool {
+// closeInfo is a map of connection id to connection metadata.
+func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient, sManager Manager, closeInfo map[string]*ConnectionCloseData) bool {
 	const op = "session.closeConnections"
 	if len(closeInfo) == 0 {
 		// This should not happen, but it's a no-op if it does. Just
@@ -496,12 +540,14 @@ func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient,
 // sessions IDs that those connections belong to. The values are
 // ignored; the parameter is expected as such just for convenience of
 // its caller.
-func makeCloseConnectionRequest(closeInfo map[string]string) *pbs.CloseConnectionRequest {
+func makeCloseConnectionRequest(closeInfo map[string]*ConnectionCloseData) *pbs.CloseConnectionRequest {
 	closeData := make([]*pbs.CloseConnectionRequestData, 0, len(closeInfo))
-	for connId := range closeInfo {
+	for connId, data := range closeInfo {
 		closeData = append(closeData, &pbs.CloseConnectionRequestData{
 			ConnectionId: connId,
 			Reason:       session.UnknownReason.String(),
+			BytesUp:      data.BytesUp,
+			BytesDown:    data.BytesDown,
 		})
 	}
 
@@ -516,7 +562,7 @@ func makeCloseConnectionRequest(closeInfo map[string]string) *pbs.CloseConnectio
 // easily lock on session once for all connections in
 // setCloseTimeForResponse.
 func makeSessionCloseInfo(
-	closeInfo map[string]string,
+	closeInfo map[string]*ConnectionCloseData,
 	response *pbs.CloseConnectionResponse,
 ) (map[string][]*pbs.CloseConnectionResponseData, error) {
 	if closeInfo == nil {
@@ -525,7 +571,8 @@ func makeSessionCloseInfo(
 
 	result := make(map[string][]*pbs.CloseConnectionResponseData)
 	for _, v := range response.GetCloseResponseData() {
-		result[closeInfo[v.GetConnectionId()]] = append(result[closeInfo[v.GetConnectionId()]], v)
+		sessionId := closeInfo[v.GetConnectionId()].SessionId
+		result[sessionId] = append(result[sessionId], v)
 	}
 
 	return result, nil
@@ -534,14 +581,15 @@ func makeSessionCloseInfo(
 // makeFakeSessionCloseInfo makes a "fake" makeFakeSessionCloseInfo, intended
 // for use when we can't contact the controller.
 func makeFakeSessionCloseInfo(
-	closeInfo map[string]string,
+	closeInfo map[string]*ConnectionCloseData,
 ) (map[string][]*pbs.CloseConnectionResponseData, error) {
 	if closeInfo == nil {
 		return nil, errMakeSessionCloseInfoNilCloseInfo
 	}
 
 	result := make(map[string][]*pbs.CloseConnectionResponseData)
-	for connectionId, sessionId := range closeInfo {
+	for connectionId, closeData := range closeInfo {
+		sessionId := closeData.SessionId
 		result[sessionId] = append(result[sessionId], &pbs.CloseConnectionResponseData{
 			ConnectionId: connectionId,
 			Status:       pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,

internal/daemon/worker/session/session_test.go

@@ -170,17 +170,26 @@ func TestSession_RequestAuthorizeConnection(t *testing.T) {
 	}
 	conn, left, err := sess.RequestAuthorizeConnection(context.Background(), "workerid", cancel)
 	require.NoError(t, err)
-	assert.Equal(t, ConnInfo{Id: "conn1", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED}, conn)
+
+	assert.Equal(t, "conn1", conn.Id)
+	assert.Equal(t, pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED, conn.Status)
+	assert.NotNil(t, conn.BytesUp)
+	assert.NotNil(t, conn.BytesDown)
+	assert.Zero(t, conn.BytesUp())
+	assert.Zero(t, conn.BytesDown())
 	assert.Equal(t, int32(-1), left)
 }
 
 func TestWorkerMakeCloseConnectionRequest(t *testing.T) {
 	require := require.New(t)
-	in := map[string]string{"foo": "one", "bar": "two"}
+	in := map[string]*ConnectionCloseData{
+		"foo": {SessionId: "one", BytesUp: 1000, BytesDown: 2000},
+		"bar": {SessionId: "two", BytesUp: 1000, BytesDown: 2000},
+	}
 	expected := &pbs.CloseConnectionRequest{
 		CloseRequestData: []*pbs.CloseConnectionRequestData{
-			{ConnectionId: "foo", Reason: session.UnknownReason.String()},
-			{ConnectionId: "bar", Reason: session.UnknownReason.String()},
+			{ConnectionId: "foo", Reason: session.UnknownReason.String(), BytesUp: 1000, BytesDown: 2000},
+			{ConnectionId: "bar", Reason: session.UnknownReason.String(), BytesUp: 1000, BytesDown: 2000},
 		},
 	}
 	actual := makeCloseConnectionRequest(in)
@@ -189,7 +198,7 @@ func TestWorkerMakeCloseConnectionRequest(t *testing.T) {
 
 func TestMakeSessionCloseInfo(t *testing.T) {
 	require := require.New(t)
-	closeInfo := map[string]string{"foo": "one", "bar": "two"}
+	closeInfo := map[string]*ConnectionCloseData{"foo": {SessionId: "one"}, "bar": {SessionId: "two"}}
 	response := &pbs.CloseConnectionResponse{
 		CloseResponseData: []*pbs.CloseConnectionResponseData{
 			{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
@@ -218,7 +227,7 @@ func TestMakeSessionCloseInfoErrorIfCloseInfoNil(t *testing.T) {
 
 func TestMakeSessionCloseInfoEmpty(t *testing.T) {
 	require := require.New(t)
-	actual, err := makeSessionCloseInfo(make(map[string]string), nil)
+	actual, err := makeSessionCloseInfo(make(map[string]*ConnectionCloseData), nil)
 	require.NoError(err)
 	require.Equal(
 		make(map[string][]*pbs.CloseConnectionResponseData),
@@ -228,7 +237,7 @@ func TestMakeSessionCloseInfoEmpty(t *testing.T) {
 
 func TestMakeFakeSessionCloseInfo(t *testing.T) {
 	require := require.New(t)
-	closeInfo := map[string]string{"foo": "one", "bar": "two"}
+	closeInfo := map[string]*ConnectionCloseData{"foo": {SessionId: "one"}, "bar": {SessionId: "two"}}
 	expected := map[string][]*pbs.CloseConnectionResponseData{
 		"one": {
 			{ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
@@ -251,10 +260,30 @@ func TestMakeFakeSessionCloseInfoErrorIfCloseInfoNil(t *testing.T) {
 
 func TestMakeFakeSessionCloseInfoEmpty(t *testing.T) {
 	require := require.New(t)
-	actual, err := makeFakeSessionCloseInfo(make(map[string]string))
+	actual, err := makeFakeSessionCloseInfo(make(map[string]*ConnectionCloseData))
 	require.NoError(err)
 	require.Equal(
 		make(map[string][]*pbs.CloseConnectionResponseData),
 		actual,
 	)
 }
+
+func TestApplyConnectionCounterCallbacks(t *testing.T) {
+	s := &sess{connInfoMap: make(map[string]*ConnInfo)}
+
+	connId := "conn1"
+	bytesUpFn := func() uint64 { return 10 }
+	bytesDnFn := func() uint64 { return 20 }
+	err := s.ApplyConnectionCounterCallbacks(connId, bytesUpFn, bytesDnFn)
+	require.EqualError(t, err, "failed to find connection info for connection id \"conn1\"")
+
+	s.connInfoMap[connId] = &ConnInfo{}
+	require.NoError(t, s.ApplyConnectionCounterCallbacks("conn1", bytesUpFn, bytesDnFn))
+
+	ci, ok := s.connInfoMap[connId]
+	require.True(t, ok)
+	require.NotNil(t, ci.BytesUp)
+	require.NotNil(t, ci.BytesDown)
+	require.EqualValues(t, ci.BytesUp(), 10)
+	require.EqualValues(t, ci.BytesDown(), 20)
+}

internal/daemon/worker/status.go

@@ -330,7 +330,7 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context, sessionManager sess
 // connections, regardless of whether or not the session is still active.
 func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionState bool, sessionManager session.Manager) {
 	const op = "worker.(Worker).cleanupConnections"
-	closeInfo := make(map[string]string)
+	closeInfo := make(map[string]*session.ConnectionCloseData)
 	cleanSessionIds := make([]string, 0)
 	sessionManager.ForEachLocalSession(func(s session.Session) bool {
 		switch {
@@ -341,8 +341,19 @@ func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionStat
 			// Cancel connections without regard to individual connection
 			// state.
 			closedIds := s.CancelAllLocalConnections()
+			localConns := s.GetLocalConnections()
 			for _, connId := range closedIds {
-				closeInfo[connId] = s.GetId()
+				var bytesUp, bytesDown uint64
+				connInfo, ok := localConns[connId]
+				if ok {
+					bytesUp = connInfo.BytesUp()
+					bytesDown = connInfo.BytesDown()
+				}
+				closeInfo[connId] = &session.ConnectionCloseData{
+					SessionId: s.GetId(),
+					BytesUp:   bytesUp,
+					BytesDown: bytesDown,
+				}
 				event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", s.GetId(), "connection_id", connId)
 			}
 
@@ -357,7 +368,7 @@ func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionStat
 			// state (ie: only ones that the controller has requested be
 			// terminated).
 			for _, connId := range s.CancelOpenLocalConnections() {
-				closeInfo[connId] = s.GetId()
+				closeInfo[connId] = &session.ConnectionCloseData{SessionId: s.GetId()}
 				event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", s.GetId(), "connection_id", connId)
 			}
 		}