Skip to content

Commit

Permalink
fix(transport): fix memory leak in grpc.Dial
Browse files Browse the repository at this point in the history
* Memoize otelgrpc.NewClientHandler for reuse

fixes: googleapis#2321
  • Loading branch information
quartzmo committed Jan 2, 2024
1 parent 9502f9c commit 98bfe5a
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 2 deletions.
24 changes: 23 additions & 1 deletion transport/grpc/dial.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import (
"net"
"os"
"strings"
"sync"
"time"

"cloud.google.com/go/compute/metadata"
Expand All @@ -27,6 +28,7 @@ import (
grpcgoogle "google.golang.org/grpc/credentials/google"
grpcinsecure "google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/credentials/oauth"
"google.golang.org/grpc/stats"

// Install grpclb, which is required for direct path.
_ "google.golang.org/grpc/balancer/grpclb"
Expand All @@ -47,6 +49,26 @@ var logRateLimiter = rate.Sometimes{Interval: 1 * time.Second}
// Assign to var for unit test replacement
var dialContext = grpc.DialContext

// otelGRPCClientHandler is a singleton otelgrpc clientHandler to be used across
// all dial connections to avoid the memory leak documented in
// https://github.com/open-telemetry/opentelemetry-go-contrib/issues/4226
//
// TODO: If 4226 has been fixed in opentelemetry-go-contrib, replace this
// singleton with inline usage for simplicity.
var (
otelGRPCClientHandler stats.Handler
once sync.Once
)

// getOtelGRPCClientHandler returns singleton otelGRPCClientHandler for reuse
// across all dial connections.
func getOtelGRPCClientHandler() stats.Handler {
once.Do(func() {
otelGRPCClientHandler = otelgrpc.NewClientHandler()
})
return otelGRPCClientHandler
}

// Dial returns a GRPC connection for use communicating with a Google cloud
// service, configured with the given ClientOptions.
func Dial(ctx context.Context, opts ...option.ClientOption) (*grpc.ClientConn, error) {
Expand Down Expand Up @@ -219,7 +241,7 @@ func addOpenTelemetryStatsHandler(opts []grpc.DialOption, settings *internal.Dia
if settings.TelemetryDisabled {
return opts
}
return append(opts, grpc.WithStatsHandler(otelgrpc.NewClientHandler()))
return append(opts, grpc.WithStatsHandler(getOtelGRPCClientHandler()))
}

// grpcTokenSource supplies PerRPCCredentials from an oauth.TokenSource.
Expand Down
2 changes: 1 addition & 1 deletion transport/grpc/dial_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ func TestLogDirectPathMisconfigNotOnGCE(t *testing.T) {
logDirectPathMisconfig(endpoint, creds.TokenSource, o)

if !metadata.OnGCE() {
wantedLog := "WARNING: DirectPath is misconfigured. DirectPath is only available in a GCE environment.."
wantedLog := "WARNING: DirectPath is misconfigured. DirectPath is only available in a GCE environment."
if !strings.Contains(buf.String(), wantedLog) {
t.Fatalf("got: %v, want: %v", buf.String(), wantedLog)
}
Expand Down

0 comments on commit 98bfe5a

Please sign in to comment.