kpt-tmr-group/pkg/grpcutil/client_conn.go

package grpcutil

import (
	"context"
	"time"

	"kpt-tmr-group/pkg/xerr"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

func DialContext(ctx context.Context, target string, options ...grpc.DialOption) (*grpc.ClientConn, error) {
	if len(options) == 0 {
		options = DefaultDialOptions()
	}

	conn, err := grpc.DialContext(ctx, target, options...)
	if err != nil {
		return nil, xerr.WithStack(err)
	}

	return conn, err
}

func DefaultDialOptions() []grpc.DialOption {
	options := DefaultAsyncDialOptions()
	return append(options, grpc.WithBlock())
}

func DefaultAsyncDialOptions() []grpc.DialOption {
	unaryInterceptors := []grpc.UnaryClientInterceptor{
		grpc_retry.UnaryClientInterceptor(
			grpc_retry.WithMax(3),
			grpc_retry.WithCodes(codes.Aborted, codes.DeadlineExceeded),
			grpc_retry.WithPerRetryTimeout(time.Millisecond*500),
		),
		grpc_prometheus.UnaryClientInterceptor,
	}
	streamInterceptors := []grpc.StreamClientInterceptor{
		grpc_retry.StreamClientInterceptor(
			grpc_retry.WithMax(3),
			grpc_retry.WithCodes(codes.Aborted, codes.DeadlineExceeded),
			grpc_retry.WithPerRetryTimeout(time.Millisecond*500),
		),
		grpc_prometheus.StreamClientInterceptor,
	}

	return []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(unaryInterceptors...)),
		grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(streamInterceptors...)),
	}
}