feat: callback notify function when connection is refused

cmd/options.go

@@ -97,3 +97,11 @@ func WithLazyRefresh() Option {
 		c.conf.LazyRefresh = true
 	}
 }
+
+// WithConnRefuseNotify configures the Proxy to call the provided function when
+// a connection is refused. The notification function is run in a goroutine.
+func WithConnRefuseNotify(n func()) Option {
+	return func(c *Command) {
+		c.connRefuseNotify = n
+	}
+}

cmd/root.go

@@ -86,10 +86,11 @@ func Execute() {
 // Command represents an invocation of the Cloud SQL Auth Proxy.
 type Command struct {
 	*cobra.Command
-	conf    *proxy.Config
-	logger  cloudsql.Logger
-	dialer  cloudsql.Dialer
-	cleanup func() error
+	conf             *proxy.Config
+	logger           cloudsql.Logger
+	dialer           cloudsql.Dialer
+	cleanup          func() error
+	connRefuseNotify func()
 }
 
 var longHelp = `
@@ -1025,7 +1026,7 @@ func runSignalWrapper(cmd *Command) (err error) {
 	startCh := make(chan *proxy.Client)
 	go func() {
 		defer close(startCh)
-		p, err := proxy.NewClient(ctx, cmd.dialer, cmd.logger, cmd.conf)
+		p, err := proxy.NewClient(ctx, cmd.dialer, cmd.logger, cmd.conf, cmd.connRefuseNotify)
 		if err != nil {
 			cmd.logger.Debugf("Error starting proxy: %v", err)
 			shutdownCh <- fmt.Errorf("unable to start: %v", err)

internal/healthcheck/healthcheck_test.go

@@ -78,7 +78,7 @@ func newProxyWithParams(t *testing.T, maxConns uint64, dialer cloudsql.Dialer, i
 		Instances:      instances,
 		MaxConnections: maxConns,
 	}
-	p, err := proxy.NewClient(context.Background(), dialer, logger, c)
+	p, err := proxy.NewClient(context.Background(), dialer, logger, c, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient: %v", err)
 	}

internal/proxy/fuse_test.go

@@ -47,7 +47,7 @@ func newTestClient(t *testing.T, d cloudsql.Dialer, fuseDir, fuseTempDir string)
 	conf := &proxy.Config{FUSEDir: fuseDir, FUSETempDir: fuseTempDir}
 
 	// This context is only used to call the Cloud SQL API
-	c, err := proxy.NewClient(context.Background(), d, testLogger, conf)
+	c, err := proxy.NewClient(context.Background(), d, testLogger, conf, nil)
 	if err != nil {
 		t.Fatalf("want error = nil, got = %v", err)
 	}
@@ -424,7 +424,7 @@ func TestFUSEWithBadDir(t *testing.T) {
 		t.Skip("skipping fuse tests in short mode.")
 	}
 	conf := &proxy.Config{FUSEDir: "/not/a/dir", FUSETempDir: randTmpDir(t)}
-	_, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, conf)
+	_, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, conf, nil)
 	if err == nil {
 		t.Fatal("proxy client should fail with bad dir")
 	}

internal/proxy/proxy.go

@@ -503,12 +503,14 @@ type Client struct {
 
 	logger cloudsql.Logger
 
+	connRefuseNotify func()
+
 	fuseMount
 }
 
 // NewClient completes the initial setup required to get the proxy to a "steady"
 // state.
-func NewClient(ctx context.Context, d cloudsql.Dialer, l cloudsql.Logger, conf *Config) (*Client, error) {
+func NewClient(ctx context.Context, d cloudsql.Dialer, l cloudsql.Logger, conf *Config, connRefuseNotify func()) (*Client, error) {
 	// Check if the caller has configured a dialer.
 	// Otherwise, initialize a new one.
 	if d == nil {
@@ -523,9 +525,10 @@ func NewClient(ctx context.Context, d cloudsql.Dialer, l cloudsql.Logger, conf *
 	}
 
 	c := &Client{
-		logger: l,
-		dialer: d,
-		conf:   conf,
+		logger:           l,
+		dialer:           d,
+		connRefuseNotify: connRefuseNotify,
+		conf:             conf,
 	}
 
 	if conf.FUSEDir != "" {
@@ -753,6 +756,9 @@ func (c *Client) serveSocketMount(_ context.Context, s *socketMount) error {
 
 			if c.conf.MaxConnections > 0 && count > c.conf.MaxConnections {
 				c.logger.Infof("max connections (%v) exceeded, refusing new connection", c.conf.MaxConnections)
+				if c.connRefuseNotify != nil {
+					go c.connRefuseNotify()
+				}
 				_ = cConn.Close()
 				return
 			}

internal/proxy/proxy_other_test.go

@@ -51,7 +51,8 @@ func TestFuseClosesGracefully(t *testing.T) {
 			FUSEDir:     t.TempDir(),
 			FUSETempDir: t.TempDir(),
 			Token:       "mytoken",
-		})
+		},
+		nil)
 	if err != nil {
 		t.Fatal(err)
 	}

internal/proxy/proxy_test.go

@@ -324,7 +324,7 @@ func TestClientInitialization(t *testing.T) {
 
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
-			c, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, tc.in)
+			c, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, tc.in, nil)
 			if err != nil {
 				t.Fatalf("want error = nil, got = %v", err)
 			}
@@ -370,7 +370,13 @@ func TestClientLimitsMaxConnections(t *testing.T) {
 		},
 		MaxConnections: 1,
 	}
-	c, err := proxy.NewClient(context.Background(), d, testLogger, in)
+	callbackGot := 0
+	connRefuseNotify := func() {
+		d.mu.Lock()
+		defer d.mu.Unlock()
+		callbackGot++
+	}
+	c, err := proxy.NewClient(context.Background(), d, testLogger, in, connRefuseNotify)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error: %v", err)
 	}
@@ -410,6 +416,10 @@ func TestClientLimitsMaxConnections(t *testing.T) {
 	if got := d.dialAttempts(); got != want {
 		t.Fatalf("dial attempts did not match expected, want = %v, got = %v", want, got)
 	}
+
+	if callbackGot == 0 {
+		t.Fatal("connRefuseNotifyCallback is not called")
+	}
 }
 
 func tryTCPDial(t *testing.T, addr string) net.Conn {
@@ -442,7 +452,7 @@ func TestClientCloseWaitsForActiveConnections(t *testing.T) {
 		},
 		WaitOnClose: 1 * time.Second,
 	}
-	c, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, in)
+	c, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error: %v", err)
 	}
@@ -464,7 +474,7 @@ func TestClientClosesCleanly(t *testing.T) {
 			{Name: "proj:reg:inst"},
 		},
 	}
-	c, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, in)
+	c, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error want = nil, got = %v", err)
 	}
@@ -486,7 +496,7 @@ func TestClosesWithError(t *testing.T) {
 			{Name: "proj:reg:inst"},
 		},
 	}
-	c, err := proxy.NewClient(context.Background(), &errorDialer{}, testLogger, in)
+	c, err := proxy.NewClient(context.Background(), &errorDialer{}, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error want = nil, got = %v", err)
 	}
@@ -542,13 +552,13 @@ func TestClientInitializationWorksRepeatedly(t *testing.T) {
 		},
 	}
 
-	c, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, in)
+	c, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("want error = nil, got = %v", err)
 	}
 	c.Close()
 
-	c, err = proxy.NewClient(ctx, &fakeDialer{}, testLogger, in)
+	c, err = proxy.NewClient(ctx, &fakeDialer{}, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("want error = nil, got = %v", err)
 	}
@@ -562,7 +572,7 @@ func TestClientNotifiesCallerOnServe(t *testing.T) {
 			{Name: "proj:region:pg"},
 		},
 	}
-	c, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, in)
+	c, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("want error = nil, got = %v", err)
 	}
@@ -595,7 +605,7 @@ func TestClientConnCount(t *testing.T) {
 		MaxConnections: 10,
 	}
 
-	c, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, in)
+	c, err := proxy.NewClient(context.Background(), &fakeDialer{}, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error: %v", err)
 	}
@@ -636,7 +646,7 @@ func TestCheckConnections(t *testing.T) {
 		},
 	}
 	d := &fakeDialer{}
-	c, err := proxy.NewClient(context.Background(), d, testLogger, in)
+	c, err := proxy.NewClient(context.Background(), d, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error: %v", err)
 	}
@@ -664,7 +674,7 @@ func TestCheckConnections(t *testing.T) {
 		},
 	}
 	ed := &errorDialer{}
-	c, err = proxy.NewClient(context.Background(), ed, testLogger, in)
+	c, err = proxy.NewClient(context.Background(), ed, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error: %v", err)
 	}
@@ -690,7 +700,7 @@ func TestRunConnectionCheck(t *testing.T) {
 		RunConnectionTest: true,
 	}
 	d := &fakeDialer{}
-	c, err := proxy.NewClient(context.Background(), d, testLogger, in)
+	c, err := proxy.NewClient(context.Background(), d, testLogger, in, nil)
 	if err != nil {
 		t.Fatalf("proxy.NewClient error: %v", err)
 	}
@@ -757,7 +767,7 @@ func TestProxyInitializationWithFailedUnixSocket(t *testing.T) {
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
-			_, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, tc.in)
+			_, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, tc.in, nil)
 			if err == nil {
 				t.Fatalf("want non nil error, got = %v", err)
 			}
@@ -801,7 +811,7 @@ func TestProxyMultiInstances(t *testing.T) {
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
-			_, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, tc.in)
+			_, err := proxy.NewClient(ctx, &fakeDialer{}, testLogger, tc.in, nil)
 			if tc.wantSuccess != (err == nil) {
 				t.Fatalf("want return = %v, got = %v", tc.wantSuccess, err == nil)
 			}