tiproxy: delete drained pods early when the connection drops to zero (#6936)

Author: YangKeao

pkg/controllers/tiproxy/tasks/finalizer.go

@@ -79,14 +79,19 @@ func drainOrDeletePod(ctx context.Context, c client.Client, state State, pod *co
 	}
 
 	remaining := time.Until(startAt.Add(time.Duration(seconds) * time.Second))
+	if remaining <= 0 {
+		return deleteTiProxyPod(ctx, c, pod)
+	}
+
+	if tiProxyConnectionsDrained(ctx, state, c, logger) {
+		return deleteTiProxyPod(ctx, c, pod)
+	}
+
 	if remaining > task.DefaultRequeueAfter {
 		remaining = task.DefaultRequeueAfter
 	}
-	if remaining > 0 {
-		return remaining, nil
-	}
 
-	return deleteTiProxyPod(ctx, c, pod)
+	return remaining, nil
 }
 
 func gracefulShutdownDeleteDelaySeconds(tiproxy *v1alpha1.TiProxy) (seconds int32, ok bool, err error) {
@@ -131,6 +136,48 @@ func deleteTiProxyPod(ctx context.Context, c client.Client, pod *corev1.Pod) (ti
 	return task.DefaultRequeueAfter, nil
 }
 
+func tiProxyConnectionsDrained(ctx context.Context, state State, c client.Client, logger logr.Logger) bool {
+	tiproxy := state.Object()
+
+	tpClient, err := newTiProxyDeleteClient(ctx, state, c)
+	if err != nil {
+		logger.Info(
+			"failed to build TiProxy API client before checking connections, continue waiting",
+			"namespace", tiproxy.Namespace,
+			"name", tiproxy.Name,
+			"error", err,
+		)
+		return false
+	}
+
+	connectionCount, err := tpClient.ConnectionCount(ctx)
+	if err != nil {
+		logger.Info(
+			"failed to query TiProxy connections before graceful delete, continue waiting",
+			"namespace", tiproxy.Namespace,
+			"name", tiproxy.Name,
+			"error", err,
+		)
+		return false
+	}
+	if connectionCount > 0 {
+		logger.Info(
+			"TiProxy still has active connections, continue waiting",
+			"namespace", tiproxy.Namespace,
+			"name", tiproxy.Name,
+			"connectionCount", connectionCount,
+		)
+		return false
+	}
+
+	logger.Info(
+		"TiProxy has no active connections, delete pod without waiting for the remaining graceful delay",
+		"namespace", tiproxy.Namespace,
+		"name", tiproxy.Name,
+	)
+	return true
+}
+
 func ensureTiProxyMarkedUnhealthy(ctx context.Context, state State, c client.Client, logger logr.Logger) bool {
 	tiproxy := state.Object()
 

pkg/controllers/tiproxy/tasks/finalizer_test.go

@@ -56,8 +56,11 @@ type testTiProxyHealthServer struct {
 	mu                  sync.Mutex
 	healthStatus        int
 	markUnhealthyStatus int
+	metricsStatus       int
+	connectionCount     int
 	healthCalls         int
 	markUnhealthyCalls  int
+	metricsCalls        int
 }
 
 func newTestTiProxyHealthServer(t *testing.T, healthStatus, markUnhealthyStatus int) *testTiProxyHealthServer {
@@ -66,6 +69,8 @@ func newTestTiProxyHealthServer(t *testing.T, healthStatus, markUnhealthyStatus
 	s := &testTiProxyHealthServer{
 		healthStatus:        healthStatus,
 		markUnhealthyStatus: markUnhealthyStatus,
+		metricsStatus:       http.StatusOK,
+		connectionCount:     1,
 	}
 
 	mux := http.NewServeMux()
@@ -84,13 +89,35 @@ func newTestTiProxyHealthServer(t *testing.T, healthStatus, markUnhealthyStatus
 		w.WriteHeader(s.healthStatus)
 		_, _ = w.Write([]byte(`{"config_checksum":1}`))
 	})
+	mux.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		s.metricsCalls++
+		w.WriteHeader(s.metricsStatus)
+		_, _ = w.Write([]byte(`# HELP tiproxy_server_connections Number of connections.
+# TYPE tiproxy_server_connections gauge
+tiproxy_server_connections ` + strconv.Itoa(s.connectionCount) + `
+`))
+	})
 	server := httptest.NewServer(mux)
 	t.Cleanup(server.Close)
 
 	s.server = server
 	return s
 }
 
+func (s *testTiProxyHealthServer) setConnectionCount(connectionCount int) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.connectionCount = connectionCount
+}
+
+func (s *testTiProxyHealthServer) setMetricsStatus(status int) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.metricsStatus = status
+}
+
 func (s *testTiProxyHealthServer) port(t *testing.T) int32 {
 	t.Helper()
 	u, err := url.Parse(s.server.URL)
@@ -107,6 +134,12 @@ func (s *testTiProxyHealthServer) counts() (health, markUnhealthy int) {
 	return s.healthCalls, s.markUnhealthyCalls
 }
 
+func (s *testTiProxyHealthServer) metricsCount() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.metricsCalls
+}
+
 func runDrainPodForDeleteTask(t *testing.T, ctx context.Context, s State, c client.Client) (task.Result, bool) {
 	t.Helper()
 
@@ -261,6 +294,55 @@ func TestTaskDrainPodForDeleteWaitForDeleteDelayAfterGracefulShutdownStarted(t *
 	assert.True(t, actual.GetDeletionTimestamp().IsZero())
 }
 
+func TestTaskDrainPodForDeleteDeletePodWhenConnectionsDropToZero(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	server := newTestTiProxyHealthServer(t, http.StatusBadGateway, http.StatusOK)
+	server.setConnectionCount(0)
+	cluster := localTestCluster()
+	tiproxy := deletingTiProxyWithAPIAddress("3600", server.port(t))
+	pod := fakePod(cluster, tiproxy)
+	pod.Annotations = map[string]string{
+		v1alpha1.AnnoKeyTiProxyGracefulShutdownBeginTime: time.Now().Add(-10 * time.Second).Format(time.RFC3339Nano),
+	}
+
+	fc := client.NewFakeClient(cluster, tiproxy, pod)
+	res, done := runDrainPodForDeleteTask(t, ctx, &state{cluster: cluster, tiproxy: tiproxy}, fc)
+
+	assert.Equal(t, task.SRetry.String(), res.Status().String())
+	assert.False(t, done)
+	assert.Equal(t, 1, server.metricsCount())
+
+	err := fc.Get(ctx, client.ObjectKeyFromObject(pod), &corev1.Pod{})
+	assert.True(t, apierrors.IsNotFound(err))
+}
+
+func TestTaskDrainPodForDeleteContinueDeleteDelayWhenConnectionMetricFails(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	server := newTestTiProxyHealthServer(t, http.StatusBadGateway, http.StatusOK)
+	server.setMetricsStatus(http.StatusInternalServerError)
+	cluster := localTestCluster()
+	tiproxy := deletingTiProxyWithAPIAddress("3600", server.port(t))
+	pod := fakePod(cluster, tiproxy)
+	pod.Annotations = map[string]string{
+		v1alpha1.AnnoKeyTiProxyGracefulShutdownBeginTime: time.Now().Add(-10 * time.Second).Format(time.RFC3339Nano),
+	}
+
+	fc := client.NewFakeClient(cluster, tiproxy, pod)
+	res, done := runDrainPodForDeleteTask(t, ctx, &state{cluster: cluster, tiproxy: tiproxy}, fc)
+
+	assert.Equal(t, task.SRetry.String(), res.Status().String())
+	assert.False(t, done)
+	assert.Equal(t, 1, server.metricsCount())
+
+	actual := &corev1.Pod{}
+	require.NoError(t, fc.Get(ctx, client.ObjectKeyFromObject(pod), actual))
+	assert.True(t, actual.GetDeletionTimestamp().IsZero())
+}
+
 func TestTaskDrainPodForDeleteDeletePodWithoutExtendingGracePeriodAfterDelay(t *testing.T) {
 	t.Parallel()
 

pkg/tiproxyapi/v1/client.go

@@ -20,18 +20,22 @@ import (
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"time"
 
 	"github.com/BurntSushi/toml"
+	"github.com/prometheus/common/expfmt"
 
 	httputil "github.com/pingcap/tidb-operator/v2/pkg/utils/http"
 )
 
 const (
-	healthPath = "api/debug/health"
-	configPath = "api/admin/config"
+	healthPath                   = "api/debug/health"
+	configPath                   = "api/admin/config"
+	metricsPath                  = "metrics"
+	tiproxyConnectionsMetricName = "tiproxy_server_connections"
 )
 
 type healthOverrideRequest struct {
@@ -45,6 +49,8 @@ type TiProxyClient interface {
 	IsHealthy(ctx context.Context) (bool, error)
 	// MarkUnhealthy makes the TiProxy health endpoint report unhealthy.
 	MarkUnhealthy(ctx context.Context) error
+	// ConnectionCount returns the current number of TiProxy SQL client connections.
+	ConnectionCount(ctx context.Context) (float64, error)
 	// SetLabels sets the labels for TiProxy.
 	SetLabels(ctx context.Context, labels map[string]string) error
 }
@@ -115,6 +121,55 @@ func (c *tiproxyClient) MarkUnhealthy(ctx context.Context) error {
 	return err
 }
 
+func (c *tiproxyClient) ConnectionCount(ctx context.Context) (float64, error) {
+	apiURL := fmt.Sprintf("%s/%s", c.url, metricsPath)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, apiURL, http.NoBody)
+	if err != nil {
+		return 0, err
+	}
+
+	//nolint:bodyclose,gosec // bodyclose: has been handled; gosec: URL is constructed from trusted internal config
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return 0, err
+	}
+	defer httputil.DeferClose(resp.Body)
+
+	if resp.StatusCode >= http.StatusBadRequest {
+		return 0, fmt.Errorf("get TiProxy metrics failed: status %d", resp.StatusCode)
+	}
+
+	return parseConnectionCount(resp.Body)
+}
+
+func parseConnectionCount(metrics io.Reader) (float64, error) {
+	parser := expfmt.TextParser{}
+	metricFamilies, err := parser.TextToMetricFamilies(metrics)
+	if err != nil {
+		return 0, fmt.Errorf("parse TiProxy metrics failed: %w", err)
+	}
+
+	mf, ok := metricFamilies[tiproxyConnectionsMetricName]
+	if !ok {
+		return 0, fmt.Errorf("metric %s not found", tiproxyConnectionsMetricName)
+	}
+
+	var count float64
+	var found bool
+	for _, metric := range mf.GetMetric() {
+		if metric.Gauge == nil {
+			continue
+		}
+		count += metric.Gauge.GetValue()
+		found = true
+	}
+	if !found {
+		return 0, fmt.Errorf("metric %s has no gauge value", tiproxyConnectionsMetricName)
+	}
+
+	return count, nil
+}
+
 func (c *tiproxyClient) SetLabels(ctx context.Context, labels map[string]string) error {
 	type labelConfig struct {
 		Labels map[string]string `toml:"labels"`

pkg/tiproxyapi/v1/client_test.go

@@ -156,3 +156,22 @@ func TestTiProxyClient_MarkUnhealthy(t *testing.T) {
 	err := client.MarkUnhealthy(context.Background())
 	require.NoError(t, err)
 }
+
+func TestTiProxyClient_ConnectionCount(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/metrics", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte(`# HELP tiproxy_server_connections Number of connections.
+# TYPE tiproxy_server_connections gauge
+tiproxy_server_connections 7
+`))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	addr := stringutil.RemoveHTTPPrefix(server.URL)
+	client := NewTiProxyClient(addr, 5*time.Second, nil)
+	count, err := client.ConnectionCount(context.Background())
+	require.NoError(t, err)
+	assert.InDelta(t, 7, count, 0)
+}