Merge branch 'master' into fix-goroutines-leak

Author: rekby

CHANGELOG.md

@@ -1,4 +1,7 @@
 * Fixed goroutines leak within topic reader on network problems
+
+## v3.67.0
+* Added `ydb.WithNodeAddressMutator` experimental option for mutate node addresses from `discovery.ListEndpoints` response
 * Added type assertion checks to enhance type safety and prevent unexpected panics in critical sections of the codebase
 
 ## v3.66.3

internal/cmd/gtrace/main.go

@@ -20,7 +20,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 )
 
-//nolint:gocyclo
+//nolint:gocyclo,funlen
 func main() {
 	var (
 		// Reports whether we were called from go:generate.

internal/cmd/gtrace/writer.go

@@ -372,6 +372,7 @@ func (w *Writer) composeHook(hook Hook, t1, t2, dst string) {
 	w.line(`}`)
 }
 
+//nolint:funlen
 func (w *Writer) composeHookCall(fn *Func, h1, h2 string) {
 	w.newScope(func() {
 		w.capture(h1, h2)

internal/credentials/oauth2.go

@@ -349,62 +349,89 @@ func (provider *oauth2TokenExchange) getRequestParams() (string, error) {
 }
 
 func (provider *oauth2TokenExchange) processTokenExchangeResponse(result *http.Response, now time.Time) error {
-	var (
-		data []byte
-		err  error
-	)
-	if result.Body != nil {
-		data, err = io.ReadAll(result.Body)
-		if err != nil {
-			return xerrors.WithStackTrace(err)
-		}
-	} else {
-		data = make([]byte, 0)
+	data, err := readResponseBody(result)
+	if err != nil {
+		return err
 	}
 
 	if result.StatusCode != http.StatusOK {
-		description := result.Status
+		return provider.handleErrorResponse(result.Status, data)
+	}
 
-		//nolint:tagliatelle
-		type errorResponse struct {
-			ErrorName        string `json:"error"`
-			ErrorDescription string `json:"error_description"`
-			ErrorURI         string `json:"error_uri"`
-		}
-		var parsedErrorResponse errorResponse
-		if err := json.Unmarshal(data, &parsedErrorResponse); err != nil {
-			description += ", could not parse response: " + err.Error()
+	parsedResponse, err := parseTokenResponse(data)
+	if err != nil {
+		return err
+	}
 
-			return xerrors.WithStackTrace(fmt.Errorf("%w: %s", errCouldNotExchangeToken, description))
-		}
+	if err := validateTokenResponse(parsedResponse, provider); err != nil {
+		return err
+	}
 
-		if parsedErrorResponse.ErrorName != "" {
-			description += ", error: " + parsedErrorResponse.ErrorName
-		}
+	provider.updateToken(parsedResponse, now)
 
-		if parsedErrorResponse.ErrorDescription != "" {
-			description += fmt.Sprintf(", description: %q", parsedErrorResponse.ErrorDescription)
-		}
+	return nil
+}
 
-		if parsedErrorResponse.ErrorURI != "" {
-			description += ", error_uri: " + parsedErrorResponse.ErrorURI
+func readResponseBody(result *http.Response) ([]byte, error) {
+	if result.Body != nil {
+		data, err := io.ReadAll(result.Body)
+		if err != nil {
+			return nil, xerrors.WithStackTrace(err)
 		}
 
-		return xerrors.WithStackTrace(fmt.Errorf("%w: %s", errCouldNotExchangeToken, description))
+		return data, nil
 	}
 
+	return make([]byte, 0), nil
+}
+
+func (provider *oauth2TokenExchange) handleErrorResponse(status string, data []byte) error {
+	description := status
+
 	//nolint:tagliatelle
-	type response struct {
-		AccessToken string `json:"access_token"`
-		TokenType   string `json:"token_type"`
-		ExpiresIn   int64  `json:"expires_in"`
-		Scope       string `json:"scope"`
+	type errorResponse struct {
+		ErrorName        string `json:"error"`
+		ErrorDescription string `json:"error_description"`
+		ErrorURI         string `json:"error_uri"`
+	}
+	var parsedErrorResponse errorResponse
+	if err := json.Unmarshal(data, &parsedErrorResponse); err != nil {
+		description += ", could not parse response: " + err.Error()
+
+		return xerrors.WithStackTrace(fmt.Errorf("%w: %s", errCouldNotExchangeToken, description))
 	}
-	var parsedResponse response
+
+	if parsedErrorResponse.ErrorName != "" {
+		description += ", error: " + parsedErrorResponse.ErrorName
+	}
+	if parsedErrorResponse.ErrorDescription != "" {
+		description += fmt.Sprintf(", description: %q", parsedErrorResponse.ErrorDescription)
+	}
+	if parsedErrorResponse.ErrorURI != "" {
+		description += ", error_uri: " + parsedErrorResponse.ErrorURI
+	}
+
+	return xerrors.WithStackTrace(fmt.Errorf("%w: %s", errCouldNotExchangeToken, description))
+}
+
+//nolint:tagliatelle
+type tokenResponse struct {
+	AccessToken string `json:"access_token"`
+	TokenType   string `json:"token_type"`
+	ExpiresIn   int64  `json:"expires_in"`
+	Scope       string `json:"scope"`
+}
+
+func parseTokenResponse(data []byte) (*tokenResponse, error) {
+	var parsedResponse tokenResponse
 	if err := json.Unmarshal(data, &parsedResponse); err != nil {
-		return xerrors.WithStackTrace(fmt.Errorf("%w: %w", errCouldNotParseResponse, err))
+		return nil, xerrors.WithStackTrace(fmt.Errorf("%w: %w", errCouldNotParseResponse, err))
 	}
 
+	return &parsedResponse, nil
+}
+
+func validateTokenResponse(parsedResponse *tokenResponse, provider *oauth2TokenExchange) error {
 	if !strings.EqualFold(parsedResponse.TokenType, "bearer") {
 		return xerrors.WithStackTrace(
 			fmt.Errorf("%w: %q", errUnsupportedTokenType, parsedResponse.TokenType))
@@ -423,18 +450,17 @@ func (provider *oauth2TokenExchange) processTokenExchangeResponse(result *http.R
 		}
 	}
 
+	return nil
+}
+
+func (provider *oauth2TokenExchange) updateToken(parsedResponse *tokenResponse, now time.Time) {
 	provider.receivedToken = "Bearer " + parsedResponse.AccessToken
 
-	// Expire time
-	expireDelta := time.Duration(parsedResponse.ExpiresIn)
-	expireDelta *= time.Second
+	expireDelta := time.Duration(parsedResponse.ExpiresIn) * time.Second
 	provider.receivedTokenExpireTime = now.Add(expireDelta)
 
-	updateDelta := time.Duration(parsedResponse.ExpiresIn / updateTimeDivider)
-	updateDelta *= time.Second
+	updateDelta := time.Duration(parsedResponse.ExpiresIn/updateTimeDivider) * time.Second
 	provider.updateTokenTime = now.Add(updateDelta)
-
-	return nil
 }
 
 func (provider *oauth2TokenExchange) exchangeToken(ctx context.Context, now time.Time) error {

internal/credentials/static.go

@@ -87,9 +87,7 @@ func (c *Static) Token(ctx context.Context) (token string, err error) {
 			fmt.Errorf("dial failed: %w", err),
 		)
 	}
-	defer func() {
-		_ = cc.Close()
-	}()
+	defer cc.Close()
 
 	client := Ydb_Auth_V1.NewAuthServiceClient(cc)
 

internal/decimal/decimal.go

@@ -224,6 +224,8 @@ func handleRemainingDigits(s string, v *big.Int, precision uint32) (*big.Int, er
 
 // Format returns the string representation of x with the given precision and
 // scale.
+//
+//nolint:funlen
 func Format(x *big.Int, precision, scale uint32) string {
 	// Check for special values and nil pointer upfront.
 	if x == nil {

internal/discovery/config/config.go

@@ -3,6 +3,8 @@ package config
 import (
 	"time"
 
+	"github.com/jonboulle/clockwork"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
@@ -15,10 +17,12 @@ const (
 type Config struct {
 	config.Common
 
-	endpoint string
-	database string
-	secure   bool
-	meta     *meta.Meta
+	endpoint       string
+	database       string
+	secure         bool
+	meta           *meta.Meta
+	addressMutator func(address string) string
+	clock          clockwork.Clock
 
 	interval time.Duration
 	trace    *trace.Discovery
@@ -28,6 +32,10 @@ func New(opts ...Option) *Config {
 	c := &Config{
 		interval: DefaultInterval,
 		trace:    &trace.Discovery{},
+		addressMutator: func(address string) string {
+			return address
+		},
+		clock: clockwork.NewRealClock(),
 	}
 	for _, opt := range opts {
 		if opt != nil {
@@ -38,10 +46,18 @@ func New(opts ...Option) *Config {
 	return c
 }
 
+func (c *Config) MutateAddress(fqdn string) string {
+	return c.addressMutator(fqdn)
+}
+
 func (c *Config) Meta() *meta.Meta {
 	return c.meta
 }
 
+func (c *Config) Clock() clockwork.Clock {
+	return c.clock
+}
+
 func (c *Config) Interval() time.Duration {
 	return c.interval
 }
@@ -85,6 +101,18 @@ func WithDatabase(database string) Option {
 	}
 }
 
+func WithClock(clock clockwork.Clock) Option {
+	return func(c *Config) {
+		c.clock = clock
+	}
+}
+
+func WithAddressMutator(addressMutator func(address string) string) Option {
+	return func(c *Config) {
+		c.addressMutator = addressMutator
+	}
+}
+
 // WithSecure set flag for secure connection
 func WithSecure(ssl bool) Option {
 	return func(c *Config) {

internal/discovery/discovery.go

@@ -19,6 +19,8 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
+//go:generate mockgen -destination grpc_client_mock_test.go -package discovery -write_package_comment=false github.com/ydb-platform/ydb-go-genproto/Ydb_Discovery_V1 DiscoveryServiceClient
+
 func New(ctx context.Context, cc grpc.ClientConnInterface, config *config.Config) *Client {
 	return &Client{
 		config: config,
@@ -35,65 +37,85 @@ type Client struct {
 	client Ydb_Discovery_V1.DiscoveryServiceClient
 }
 
-// Discover cluster endpoints
-func (c *Client) Discover(ctx context.Context) (endpoints []endpoint.Endpoint, err error) {
+func discover(
+	ctx context.Context,
+	client Ydb_Discovery_V1.DiscoveryServiceClient,
+	config *config.Config,
+) (endpoints []endpoint.Endpoint, location string, err error) {
 	var (
-		onDone = trace.DiscoveryOnDiscover(
-			c.config.Trace(), &ctx,
-			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/discovery.(*Client).Discover"),
-			c.config.Endpoint(), c.config.Database(),
-		)
 		request = Ydb_Discovery.ListEndpointsRequest{
-			Database: c.config.Database(),
+			Database: config.Database(),
 		}
 		response *Ydb_Discovery.ListEndpointsResponse
 		result   Ydb_Discovery.ListEndpointsResult
-		location string
 	)
-	defer func() {
-		nodes := make([]trace.EndpointInfo, 0, len(endpoints))
-		for _, e := range endpoints {
-			nodes = append(nodes, e.Copy())
-		}
-		onDone(location, nodes, err)
-	}()
-
-	ctx, err = c.config.Meta().Context(ctx)
-	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
-	}
 
-	response, err = c.client.ListEndpoints(ctx, &request)
+	response, err = client.ListEndpoints(ctx, &request)
 	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
+		return nil, location, xerrors.WithStackTrace(err)
 	}
 
 	if response.GetOperation().GetStatus() != Ydb.StatusIds_SUCCESS {
-		return nil, xerrors.WithStackTrace(
+		return nil, location, xerrors.WithStackTrace(
 			xerrors.FromOperation(response.GetOperation()),
 		)
 	}
 
 	err = response.GetOperation().GetResult().UnmarshalTo(&result)
 	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
+		return nil, location, xerrors.WithStackTrace(err)
 	}
 
 	location = result.GetSelfLocation()
 	endpoints = make([]endpoint.Endpoint, 0, len(result.GetEndpoints()))
 	for _, e := range result.GetEndpoints() {
-		if e.GetSsl() == c.config.Secure() {
+		if e.GetSsl() == config.Secure() {
 			endpoints = append(endpoints, endpoint.New(
-				net.JoinHostPort(e.GetAddress(), strconv.Itoa(int(e.GetPort()))),
+				net.JoinHostPort(
+					config.MutateAddress(e.GetAddress()),
+					strconv.Itoa(int(e.GetPort())),
+				),
 				endpoint.WithLocation(e.GetLocation()),
 				endpoint.WithID(e.GetNodeId()),
 				endpoint.WithLoadFactor(e.GetLoadFactor()),
 				endpoint.WithLocalDC(e.GetLocation() == location),
 				endpoint.WithServices(e.GetService()),
+				endpoint.WithLastUpdated(config.Clock().Now()),
 			))
 		}
 	}
 
+	return endpoints, result.GetSelfLocation(), nil
+}
+
+// Discover cluster endpoints
+func (c *Client) Discover(ctx context.Context) (endpoints []endpoint.Endpoint, finalErr error) {
+	var (
+		onDone = trace.DiscoveryOnDiscover(
+			c.config.Trace(), &ctx,
+			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/discovery.(*Client).Discover"),
+			c.config.Endpoint(), c.config.Database(),
+		)
+		location string
+	)
+	defer func() {
+		nodes := make([]trace.EndpointInfo, 0, len(endpoints))
+		for _, e := range endpoints {
+			nodes = append(nodes, e.Copy())
+		}
+		onDone(location, nodes, finalErr)
+	}()
+
+	ctx, err := c.config.Meta().Context(ctx)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
+
+	endpoints, location, err = discover(ctx, c.client, c.config)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
+
 	return endpoints, nil
 }