agent: smarter readiness check

This commit introduces a smarter way for the readiness check.

Currently as soon as the agent is up, it signals itself as ready
regardless of whether its actually connected to the server or not.

This commit introduces a readiness check in the Agent which reports
true if its connected to atleast one proxy server.

Signed-off-by: Imran Pochi <[email protected]>

Author: ipochi

cmd/agent/app/server.go

@@ -36,6 +36,7 @@ import (
 	"k8s.io/klog/v2"
 
 	"sigs.k8s.io/apiserver-network-proxy/cmd/agent/app/options"
+	"sigs.k8s.io/apiserver-network-proxy/pkg/agent"
 	"sigs.k8s.io/apiserver-network-proxy/pkg/util"
 )
 
@@ -63,11 +64,13 @@ func (a *Agent) run(o *options.GrpcProxyAgentOptions) error {
 	}
 
 	stopCh := make(chan struct{})
-	if err := a.runProxyConnection(o, stopCh); err != nil {
+
+	cs, err := a.runProxyConnection(o, stopCh)
+	if err != nil {
 		return fmt.Errorf("failed to run proxy connection with %v", err)
 	}
 
-	if err := a.runHealthServer(o); err != nil {
+	if err := a.runHealthServer(o, cs); err != nil {
 		return fmt.Errorf("failed to run health server with %v", err)
 	}
 
@@ -80,11 +83,11 @@ func (a *Agent) run(o *options.GrpcProxyAgentOptions) error {
 	return nil
 }
 
-func (a *Agent) runProxyConnection(o *options.GrpcProxyAgentOptions, stopCh <-chan struct{}) error {
+func (a *Agent) runProxyConnection(o *options.GrpcProxyAgentOptions, stopCh <-chan struct{}) (agent.ReadinessManager, error) {
 	var tlsConfig *tls.Config
 	var err error
 	if tlsConfig, err = util.GetClientTLSConfig(o.CaCert, o.AgentCert, o.AgentKey, o.ProxyServerHost, o.AlpnProtos); err != nil {
-		return err
+		return nil, err
 	}
 	dialOptions := []grpc.DialOption{
 		grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)),
@@ -97,15 +100,21 @@ func (a *Agent) runProxyConnection(o *options.GrpcProxyAgentOptions, stopCh <-ch
 	cs := cc.NewAgentClientSet(stopCh)
 	cs.Serve()
 
-	return nil
+	return cs, nil
 }
 
-func (a *Agent) runHealthServer(o *options.GrpcProxyAgentOptions) error {
+func (a *Agent) runHealthServer(o *options.GrpcProxyAgentOptions, cs agent.ReadinessManager) error {
 	livenessHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		fmt.Fprintf(w, "ok")
 	})
 	readinessHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		fmt.Fprintf(w, "ok")
+		if cs.Ready() {
+			w.WriteHeader(http.StatusOK)
+			fmt.Fprintf(w, "ok")
+		} else {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			fmt.Fprintf(w, "not ready")
+		}
 	})
 
 	muxHandler := http.NewServeMux()

pkg/agent/readiness.go

@@ -0,0 +1,30 @@
+/*
+copyright 2023 the kubernetes authors.
+
+licensed under the apache license, version 2.0 (the "license");
+you may not use this file except in compliance with the license.
+you may obtain a copy of the license at
+
+    http://www.apache.org/licenses/license-2.0
+
+unless required by applicable law or agreed to in writing, software
+distributed under the license is distributed on an "as is" basis,
+without warranties or conditions of any kind, either express or implied.
+see the license for the specific language governing permissions and
+limitations under the license.
+*/
+
+package agent
+
+// ReadinessManager supports checking if the agent is ready.
+type ReadinessManager interface {
+	// Ready returns true the proxy server is ready.
+	Ready() bool
+}
+
+var _ ReadinessManager = &ClientSet{}
+
+func (cs *ClientSet) Ready() bool {
+	// Returns true if the agent is connected to at least one server.
+	return cs.HealthyClientsCount() > 0
+}

tests/proxy_test.go

@@ -969,6 +969,7 @@ func runAgentWithID(agentID, addr string, stopCh <-chan struct{}) *agent.ClientS
 		ProbeInterval: 100 * time.Millisecond,
 		DialOptions:   []grpc.DialOption{grpc.WithInsecure()},
 	}
+
 	client := cc.NewAgentClientSet(stopCh)
 	client.Serve()
 	return client

tests/readiness_test.go

@@ -16,11 +16,9 @@ limitations under the License.
 
 package tests
 
-import (
-	"testing"
-)
+import "testing"
 
-func TestReadiness(t *testing.T) {
+func TestGRPCServerAndAgentReadiness(t *testing.T) {
 	stopCh := make(chan struct{})
 	defer close(stopCh)
 
@@ -38,8 +36,45 @@ func TestReadiness(t *testing.T) {
 	clientset := runAgent(proxy.agent, stopCh)
 	waitForConnectedServerCount(t, 1, clientset)
 
+	// check the agent connected status
+	isAgentReady := clientset.Ready()
+	if !isAgentReady {
+		t.Fatalf("expected connection status 'true', got: %t", isAgentReady)
+	}
 	ready, _ = server.Readiness.Ready()
 	if !ready {
 		t.Fatalf("expected ready")
 	}
 }
+
+func TestHTTPConnServerAndAgentReadiness(t *testing.T) {
+	stopCh := make(chan struct{})
+	defer close(stopCh)
+
+	proxy, cleanup, err := runHTTPConnProxyServer()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cleanup()
+
+	clientset := runAgent(proxy.agent, stopCh)
+	waitForConnectedServerCount(t, 1, clientset)
+
+	// check the agent connected status
+	isAgentReady := clientset.Ready()
+	if !isAgentReady {
+		t.Fatalf("expected connection status 'true', got: %t", isAgentReady)
+	}
+}
+
+func TestAgentReadinessWithoutServer(t *testing.T) {
+	stopCh := make(chan struct{})
+	defer close(stopCh)
+
+	clientset := runAgent("localhost:8080", stopCh)
+	// check the agent connected status
+	isAgentReady := clientset.Ready()
+	if isAgentReady {
+		t.Fatalf("expected connection status 'false', got: %t", isAgentReady)
+	}
+}