Merge pull request kubernetes-sigs#599 from jkh52/backend-tests

Add test coverage around BackendManager and BackendStorage.
jkh52 · Mar 31, 2024 · f97c7af · f97c7af
2 parents 80db82a + d1fd7dc
commit f97c7af
Show file tree

Hide file tree

Showing 7 changed files with 585 additions and 109 deletions.
diff --git a/Makefile b/Makefile
@@ -64,6 +64,13 @@ mock_gen:
 	cat hack/go-license-header.txt proto/agent/mocks/agent_mock.go > proto/agent/mocks/agent_mock.licensed.go
 	mv proto/agent/mocks/agent_mock.licensed.go proto/agent/mocks/agent_mock.go
 
+# Unit tests with faster execution (nicer for development).
+.PHONY: fast-test
+fast-test:
+	go test -mod=vendor -race ./...
+	cd konnectivity-client && go test -race ./...
+
+# Unit tests with fuller coverage, invoked by CI system.
 .PHONY: test
 test:
 	go test -mod=vendor -race -covermode=atomic -coverprofile=konnectivity.out ./... && go tool cover -html=konnectivity.out -o=konnectivity.html

diff --git a/pkg/server/backend_manager.go b/pkg/server/backend_manager.go
@@ -182,10 +182,10 @@ func NewBackend(conn agent.AgentService_ConnectServer) (Backend, error) {
 // BackendStorage is an interface to manage the storage of the backend
 // connections, i.e., get, add and remove
 type BackendStorage interface {
-	// AddBackend adds a backend.
-	AddBackend(identifier string, idType header.IdentifierType, backend Backend)
-	// RemoveBackend removes a backend.
-	RemoveBackend(identifier string, idType header.IdentifierType, backend Backend)
+	// addBackend adds a backend.
+	addBackend(identifier string, idType header.IdentifierType, backend Backend)
+	// removeBackend removes a backend.
+	removeBackend(identifier string, idType header.IdentifierType, backend Backend)
 	// NumBackends returns the number of backends.
 	NumBackends() int
 }
@@ -199,6 +199,10 @@ type BackendManager interface {
 	// pick a backend for every tunnel session and each tunnel session may
 	// contains multiple requests.
 	Backend(ctx context.Context) (Backend, error)
+	// AddBackend adds a backend.
+	AddBackend(backend Backend)
+	// RemoveBackend adds a backend.
+	RemoveBackend(backend Backend)
 	BackendStorage
 	ReadinessManager
 }
@@ -215,13 +219,28 @@ func (dbm *DefaultBackendManager) Backend(_ context.Context) (Backend, error) {
 	return dbm.DefaultBackendStorage.GetRandomBackend()
 }
 
+func (dbm *DefaultBackendManager) AddBackend(backend Backend) {
+	agentID := backend.GetAgentID()
+	klog.V(5).InfoS("Add the agent to DefaultBackendManager", "agentID", agentID)
+	dbm.addBackend(agentID, header.UID, backend)
+}
+
+func (dbm *DefaultBackendManager) RemoveBackend(backend Backend) {
+	agentID := backend.GetAgentID()
+	klog.V(5).InfoS("Remove the agent from the DefaultBackendManager", "agentID", agentID)
+	dbm.removeBackend(agentID, header.UID, backend)
+}
+
 // DefaultBackendStorage is the default backend storage.
 type DefaultBackendStorage struct {
 	mu sync.RWMutex //protects the following
 	// A map between agentID and its grpc connections.
 	// 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.
+	//
+	// TODO: fix documentation. This is not always agentID, e.g. in
+	// the case of DestHostBackendManager.
 	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
@@ -267,8 +286,8 @@ func containIDType(idTypes []header.IdentifierType, idType header.IdentifierType
 	return false
 }
 
-// AddBackend adds a backend.
-func (s *DefaultBackendStorage) AddBackend(identifier string, idType header.IdentifierType, backend Backend) {
+// addBackend adds a backend.
+func (s *DefaultBackendStorage) addBackend(identifier string, idType header.IdentifierType, backend Backend) {
 	if !containIDType(s.idTypes, idType) {
 		klog.V(4).InfoS("fail to add backend", "backend", identifier, "error", &ErrWrongIDType{idType, s.idTypes})
 		return
@@ -295,8 +314,8 @@ func (s *DefaultBackendStorage) AddBackend(identifier string, idType header.Iden
 	}
 }
 
-// RemoveBackend removes a backend.
-func (s *DefaultBackendStorage) RemoveBackend(identifier string, idType header.IdentifierType, backend Backend) {
+// removeBackend removes a backend.
+func (s *DefaultBackendStorage) removeBackend(identifier string, idType header.IdentifierType, backend Backend) {
 	if !containIDType(s.idTypes, idType) {
 		klog.ErrorS(&ErrWrongIDType{idType, s.idTypes}, "fail to remove backend")
 		return