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

Author: bhshkh

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

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=

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

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.

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,
 			}))