From acd258182eaf544256fba965e079b4198ad41efc Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Mon, 2 Dec 2024 09:49:15 -0800 Subject: [PATCH] feat: Async refresh dry run in parallel with sync refresh (#2849) * feat: Add async refresh dry run * add meter provider and attributes * adding error handler instead of metrics * Rename variable and add comment * adding more comments * Updating oauth2adapt --- go.mod | 2 +- go.sum | 4 +-- internal/settings.go | 3 ++ option/internaloption/internaloption.go | 27 ++++++++++++++++ transport/grpc/dial.go | 42 ++++++++++++++++++++++++- 5 files changed, 74 insertions(+), 4 deletions(-) diff --git a/go.mod b/go.mod index a87601462f..80a1dbeb94 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.21 require ( cloud.google.com/go v0.116.0 cloud.google.com/go/auth v0.10.2 - cloud.google.com/go/auth/oauth2adapt v0.2.5 + cloud.google.com/go/auth/oauth2adapt v0.2.6 cloud.google.com/go/compute/metadata v0.5.2 github.com/google/go-cmp v0.6.0 github.com/google/s2a-go v0.1.8 diff --git a/go.sum b/go.sum index 76eea2615b..8c39f32033 100644 --- a/go.sum +++ b/go.sum @@ -3,8 +3,8 @@ cloud.google.com/go v0.116.0 h1:B3fRrSDkLRt5qSHWe40ERJvhvnQwdZiHu0bJOpldweE= cloud.google.com/go v0.116.0/go.mod h1:cEPSRWPzZEswwdr9BxE6ChEn01dWlTaF05LiC2Xs70U= cloud.google.com/go/auth v0.10.2 h1:oKF7rgBfSHdp/kuhXtqU/tNDr0mZqhYbEh+6SiqzkKo= cloud.google.com/go/auth v0.10.2/go.mod h1:xxA5AqpDrvS+Gkmo9RqrGGRh6WSNKKOXhY3zNOr38tI= -cloud.google.com/go/auth/oauth2adapt v0.2.5 h1:2p29+dePqsCHPP1bqDJcKj4qxRyYCcbzKpFyKGt3MTk= -cloud.google.com/go/auth/oauth2adapt v0.2.5/go.mod h1:AlmsELtlEBnaNTL7jCj8VQFLy6mbZv0s4Q7NGBeQ5E8= +cloud.google.com/go/auth/oauth2adapt v0.2.6 h1:V6a6XDu2lTwPZWOawrAa9HUK+DB2zfJyTuciBG5hFkU= +cloud.google.com/go/auth/oauth2adapt v0.2.6/go.mod h1:AlmsELtlEBnaNTL7jCj8VQFLy6mbZv0s4Q7NGBeQ5E8= cloud.google.com/go/compute/metadata v0.5.2 h1:UxK4uu/Tn+I3p2dYWTfiX4wva7aYlKixAHn3fyqngqo= cloud.google.com/go/compute/metadata v0.5.2/go.mod h1:C66sj2AluDcIqakBq/M8lw8/ybHgOZqin2obFxa/E5k= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= diff --git a/internal/settings.go b/internal/settings.go index 4839ab6d4f..4f5b1a0ebe 100644 --- a/internal/settings.go +++ b/internal/settings.go @@ -72,6 +72,9 @@ type DialSettings struct { // New Auth library Options AuthCredentials *auth.Credentials EnableNewAuthLibrary bool + + // TODO(b/372244283): Remove after b/358175516 has been fixed + EnableAsyncRefreshDryRun func() } // GetScopes returns the user-provided scopes, if set, or else falls back to the diff --git a/option/internaloption/internaloption.go b/option/internaloption/internaloption.go index 209045ba12..c63c0c194a 100644 --- a/option/internaloption/internaloption.go +++ b/option/internaloption/internaloption.go @@ -211,6 +211,33 @@ func (w enableNewAuthLibrary) Apply(o *internal.DialSettings) { o.EnableNewAuthLibrary = bool(w) } +// EnableAsyncRefreshDryRun returns a ClientOption that specifies if libraries in this +// module should asynchronously refresh auth token in parallel to sync refresh. +// +// This option can be used to determine whether refreshing the token asymnchronously +// prior to its actual expiry works without any issues in a particular environment. +// +// errHandler function will be called when there is an error while refreshing +// the token asynchronously. +// +// This is an EXPERIMENTAL option and will be removed in the future. +// TODO(b/372244283): Remove after b/358175516 has been fixed +func EnableAsyncRefreshDryRun(errHandler func()) option.ClientOption { + return enableAsyncRefreshDryRun{ + errHandler: errHandler, + } +} + +// TODO(b/372244283): Remove after b/358175516 has been fixed +type enableAsyncRefreshDryRun struct { + errHandler func() +} + +// TODO(b/372244283): Remove after b/358175516 has been fixed +func (w enableAsyncRefreshDryRun) Apply(o *internal.DialSettings) { + o.EnableAsyncRefreshDryRun = w.errHandler +} + // EmbeddableAdapter is a no-op option.ClientOption that allow libraries to // create their own client options by embedding this type into their own // client-specific option wrapper. See example for usage. diff --git a/transport/grpc/dial.go b/transport/grpc/dial.go index 22940678d2..4015df6303 100644 --- a/transport/grpc/dial.go +++ b/transport/grpc/dial.go @@ -262,6 +262,40 @@ func prepareDialOptsNewAuth(ds *internal.DialSettings) []grpc.DialOption { return append(opts, ds.GRPCDialOpts...) } +// dryRunAsync is a wrapper for oauth2.TokenSource that performs a sync refresh +// after an async refresh. Token generated by async refresh is not used. +// +// This is an EXPERIMENTAL feature and may be removed or changed in the future. +// It is a temporary struct to determine if the async refresh +// is working properly. +// TODO(b/372244283): Remove after b/358175516 has been fixed +type dryRunAsync struct { + asyncTokenSource oauth2.TokenSource + syncTokenSource oauth2.TokenSource + errHandler func() +} + +// TODO(b/372244283): Remove after b/358175516 has been fixed +func newDryRunAsync(ts oauth2.TokenSource, errHandler func()) dryRunAsync { + tp := auth.NewCachedTokenProvider(oauth2adapt.TokenProviderFromTokenSource(ts), nil) + asyncTs := oauth2adapt.TokenSourceFromTokenProvider(tp) + return dryRunAsync{ + syncTokenSource: ts, + asyncTokenSource: asyncTs, + errHandler: errHandler, + } +} + +// Token returns a token or an error. +// TODO(b/372244283): Remove after b/358175516 has been fixed +func (async dryRunAsync) Token() (*oauth2.Token, error) { + _, err := async.asyncTokenSource.Token() + if err != nil { + async.errHandler() + } + return async.syncTokenSource.Token() +} + func dial(ctx context.Context, insecure bool, o *internal.DialSettings) (*grpc.ClientConn, error) { if o.HTTPClient != nil { return nil, errors.New("unsupported HTTP client specified") @@ -298,8 +332,14 @@ func dial(ctx context.Context, insecure bool, o *internal.DialSettings) (*grpc.C if err != nil { return nil, err } + + ts := creds.TokenSource + // TODO(b/372244283): Remove after b/358175516 has been fixed + if o.EnableAsyncRefreshDryRun != nil { + ts = newDryRunAsync(ts, o.EnableAsyncRefreshDryRun) + } grpcOpts = append(grpcOpts, grpc.WithPerRPCCredentials(grpcTokenSource{ - TokenSource: oauth.TokenSource{TokenSource: creds.TokenSource}, + TokenSource: oauth.TokenSource{TokenSource: ts}, quotaProject: internal.GetQuotaProject(creds, o.QuotaProject), requestReason: o.RequestReason, }))