alts: Read max number of concurrent ALTS handshakes from environment variable.

credentials/alts/internal/handshaker/handshaker.go

@@ -25,6 +25,8 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"os"
+	"strconv"
 	"sync"
 
 	grpc "google.golang.org/grpc"
@@ -41,9 +43,10 @@ const (
 	// The maximum byte size of receive frames.
 	frameLimit              = 64 * 1024 // 64 KB
 	rekeyRecordProtocolName = "ALTSRP_GCM_AES128_REKEY"
-	// maxPendingHandshakes represents the maximum number of concurrent
+	// defaultMaxPendingHandshakes represents the default maximum number of concurrent
 	// handshakes.
-	maxPendingHandshakes = 100
+	defaultMaxPendingHandshakes        = 100
+	maxConcurrentHandshakesEnvVariable = "GRPC_ALTS_MAX_CONCURRENT_HANDSHAKES"
 )
 
 var (
@@ -59,6 +62,7 @@ var (
 			return conn.NewAES128GCMRekey(s, keyData)
 		},
 	}
+	maxPendingHandshakes int64
 	// control number of concurrent created (but not closed) handshakers.
 	mu                   sync.Mutex
 	concurrentHandshakes = int64(0)
@@ -70,6 +74,7 @@ var (
 )
 
 func init() {
+	maxPendingHandshakes = getMaxConcurrentHandshakes()
 	for protocol, f := range altsRecordFuncs {
 		if err := conn.RegisterProtocol(protocol, f); err != nil {
 			panic(err)
@@ -391,3 +396,18 @@ func (h *altsHandshaker) Close() {
 		h.stream.CloseSend()
 	}
 }
+
+func getMaxConcurrentHandshakes() int64 {
+	maxConcurrentHandshakesString, ok := os.LookupEnv(maxConcurrentHandshakesEnvVariable)
+	if !ok {
+		return defaultMaxPendingHandshakes
+	}
+	maxConcurrentHandshakes, err := strconv.ParseInt(maxConcurrentHandshakesString, 10, 64)
+	if err != nil {
+		return defaultMaxPendingHandshakes
+	}
+	if maxConcurrentHandshakes < 0 {
+		return defaultMaxPendingHandshakes
+	}
+	return maxConcurrentHandshakes
+}

credentials/alts/internal/handshaker/handshaker_test.go

@@ -22,6 +22,7 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"os"
 	"testing"
 	"time"
 
@@ -134,7 +135,7 @@ func (s) TestClientHandshake(t *testing.T) {
 		numberOfHandshakes int
 	}{
 		{0 * time.Millisecond, 1},
-		{100 * time.Millisecond, 10 * maxPendingHandshakes},
+		{100 * time.Millisecond, 10 * int(maxPendingHandshakes)},
 	} {
 		errc := make(chan error)
 		stat.Reset()
@@ -182,7 +183,7 @@ func (s) TestClientHandshake(t *testing.T) {
 		}
 
 		// Ensure that there are no concurrent calls more than the limit.
-		if stat.MaxConcurrentCalls > maxPendingHandshakes {
+		if stat.MaxConcurrentCalls > int(maxPendingHandshakes) {
 			t.Errorf("Observed %d concurrent handshakes; want <= %d", stat.MaxConcurrentCalls, maxPendingHandshakes)
 		}
 	}
@@ -194,7 +195,7 @@ func (s) TestServerHandshake(t *testing.T) {
 		numberOfHandshakes int
 	}{
 		{0 * time.Millisecond, 1},
-		{100 * time.Millisecond, 10 * maxPendingHandshakes},
+		{100 * time.Millisecond, 10 * int(maxPendingHandshakes)},
 	} {
 		errc := make(chan error)
 		stat.Reset()
@@ -239,7 +240,7 @@ func (s) TestServerHandshake(t *testing.T) {
 		}
 
 		// Ensure that there are no concurrent calls more than the limit.
-		if stat.MaxConcurrentCalls > maxPendingHandshakes {
+		if stat.MaxConcurrentCalls > int(maxPendingHandshakes) {
 			t.Errorf("Observed %d concurrent handshakes; want <= %d", stat.MaxConcurrentCalls, maxPendingHandshakes)
 		}
 	}
@@ -349,3 +350,25 @@ func (s) TestNewServerHandshaker(t *testing.T) {
 	}
 	hs.Close()
 }
+
+func (s) TestGetMaxConcurrentHandshakes(t *testing.T) {
+	for _, testCase := range []struct {
+		setMaxConcurrentHandshakeEnvVariable bool
+		maxConcurrentHandshakeString         string
+		expectedMaxConcurrentHandshakes      int64
+	}{
+		{false, "", defaultMaxPendingHandshakes},
+		{true, "", defaultMaxPendingHandshakes},
+		{true, "not-an-integer", defaultMaxPendingHandshakes},
+		{true, "-10", defaultMaxPendingHandshakes},
+		{true, "10", int64(10)},
+	} {
+		os.Unsetenv(maxConcurrentHandshakesEnvVariable)
+		if testCase.setMaxConcurrentHandshakeEnvVariable {
+			os.Setenv(maxConcurrentHandshakesEnvVariable, testCase.maxConcurrentHandshakeString)
+		}
+		if got, want := getMaxConcurrentHandshakes(), testCase.expectedMaxConcurrentHandshakes; got != want {
+			t.Errorf("getMaxConcurrentHandshakes() = %d, want = %d", got, want)
+		}
+	}
+}