2020-12-02 15:09:37 +00:00
|
|
|
package gitaly
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"strings"
|
|
|
|
"sync"
|
|
|
|
|
2021-07-21 15:08:52 +00:00
|
|
|
"github.com/golang/protobuf/jsonpb" //lint:ignore SA1019 https://gitlab.com/gitlab-org/gitlab/-/issues/324868
|
|
|
|
"github.com/golang/protobuf/proto" //lint:ignore SA1019 https://gitlab.com/gitlab-org/gitlab/-/issues/324868
|
2020-12-02 15:09:37 +00:00
|
|
|
grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
|
|
|
|
grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
2020-12-09 03:09:24 +00:00
|
|
|
"github.com/prometheus/client_golang/prometheus/promauto"
|
2021-10-11 09:09:08 +00:00
|
|
|
"github.com/sirupsen/logrus"
|
2021-07-21 15:08:52 +00:00
|
|
|
"google.golang.org/grpc"
|
|
|
|
"google.golang.org/grpc/metadata"
|
|
|
|
|
2022-06-30 09:10:05 +00:00
|
|
|
gitalyauth "gitlab.com/gitlab-org/gitaly/v15/auth"
|
|
|
|
gitalyclient "gitlab.com/gitlab-org/gitaly/v15/client"
|
|
|
|
"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
|
2020-12-02 15:09:37 +00:00
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
"gitlab.com/gitlab-org/gitlab/workhorse/internal/api"
|
|
|
|
|
2020-12-02 15:09:37 +00:00
|
|
|
grpccorrelation "gitlab.com/gitlab-org/labkit/correlation/grpc"
|
|
|
|
grpctracing "gitlab.com/gitlab-org/labkit/tracing/grpc"
|
|
|
|
)
|
|
|
|
|
2021-10-11 09:09:08 +00:00
|
|
|
type cacheKey struct {
|
|
|
|
address, token string
|
|
|
|
}
|
2020-12-02 15:09:37 +00:00
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
func getCacheKey(server api.GitalyServer) cacheKey {
|
2022-04-28 12:09:06 +00:00
|
|
|
return cacheKey{address: server.Address, token: server.Token}
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
type connectionsCache struct {
|
|
|
|
sync.RWMutex
|
|
|
|
connections map[cacheKey]*grpc.ClientConn
|
|
|
|
}
|
|
|
|
|
|
|
|
var (
|
|
|
|
jsonUnMarshaler = jsonpb.Unmarshaler{AllowUnknownFields: true}
|
2021-10-11 09:09:08 +00:00
|
|
|
// This connection cache map contains two types of connections:
|
|
|
|
// - Normal gRPC connections
|
|
|
|
// - Sidechannel connections. When client dials to the Gitaly server, the
|
|
|
|
// server multiplexes the connection using Yamux. In the future, the server
|
|
|
|
// can open another stream to transfer data without gRPC. Besides, we apply
|
|
|
|
// a framing protocol to add the half-close capability to Yamux streams.
|
|
|
|
// Hence, we cannot use those connections interchangeably.
|
|
|
|
cache = connectionsCache{
|
2020-12-02 15:09:37 +00:00
|
|
|
connections: make(map[cacheKey]*grpc.ClientConn),
|
|
|
|
}
|
2021-10-11 09:09:08 +00:00
|
|
|
sidechannelRegistry *gitalyclient.SidechannelRegistry
|
2020-12-02 15:09:37 +00:00
|
|
|
|
2020-12-09 03:09:24 +00:00
|
|
|
connectionsTotal = promauto.NewCounterVec(
|
2020-12-02 15:09:37 +00:00
|
|
|
prometheus.CounterOpts{
|
|
|
|
Name: "gitlab_workhorse_gitaly_connections_total",
|
|
|
|
Help: "Number of Gitaly connections that have been established",
|
|
|
|
},
|
|
|
|
[]string{"status"},
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
2021-10-11 09:09:08 +00:00
|
|
|
func InitializeSidechannelRegistry(logger *logrus.Logger) {
|
|
|
|
if sidechannelRegistry == nil {
|
|
|
|
sidechannelRegistry = gitalyclient.NewSidechannelRegistry(logrus.NewEntry(logger))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
type MetadataFunc func(metadata.MD)
|
|
|
|
|
|
|
|
func WithUserID(userID string) MetadataFunc {
|
|
|
|
return func(md metadata.MD) {
|
|
|
|
md.Append("user_id", userID)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithUsername(username string) MetadataFunc {
|
|
|
|
return func(md metadata.MD) {
|
|
|
|
md.Append("username", username)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithFeatures(features map[string]string) MetadataFunc {
|
|
|
|
return func(md metadata.MD) {
|
|
|
|
for k, v := range features {
|
|
|
|
if !strings.HasPrefix(k, "gitaly-feature-") {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
md.Append(k, v)
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
2022-05-18 09:07:58 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func withOutgoingMetadata(ctx context.Context, addMetadataFuncs ...MetadataFunc) context.Context {
|
|
|
|
md := metadata.New(nil)
|
|
|
|
|
|
|
|
for _, f := range addMetadataFuncs {
|
|
|
|
f(md)
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
return metadata.NewOutgoingContext(ctx, md)
|
|
|
|
}
|
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
func NewSmartHTTPClient(ctx context.Context, server api.GitalyServer, metadataFuncs ...MetadataFunc) (context.Context, *SmartHTTPClient, error) {
|
2020-12-02 15:09:37 +00:00
|
|
|
conn, err := getOrCreateConnection(server)
|
|
|
|
if err != nil {
|
|
|
|
return nil, nil, err
|
|
|
|
}
|
|
|
|
grpcClient := gitalypb.NewSmartHTTPServiceClient(conn)
|
2021-10-11 09:09:08 +00:00
|
|
|
smartHTTPClient := &SmartHTTPClient{
|
|
|
|
SmartHTTPServiceClient: grpcClient,
|
|
|
|
sidechannelRegistry: sidechannelRegistry,
|
|
|
|
}
|
2022-05-18 09:07:58 +00:00
|
|
|
|
|
|
|
return withOutgoingMetadata(
|
|
|
|
ctx,
|
|
|
|
metadataFuncs...,
|
|
|
|
), smartHTTPClient, nil
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
func NewBlobClient(ctx context.Context, server api.GitalyServer, addMetadataFuncs ...MetadataFunc) (context.Context, *BlobClient, error) {
|
2020-12-02 15:09:37 +00:00
|
|
|
conn, err := getOrCreateConnection(server)
|
|
|
|
if err != nil {
|
|
|
|
return nil, nil, err
|
|
|
|
}
|
|
|
|
grpcClient := gitalypb.NewBlobServiceClient(conn)
|
2022-05-18 09:07:58 +00:00
|
|
|
return withOutgoingMetadata(ctx, addMetadataFuncs...), &BlobClient{grpcClient}, nil
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
func NewRepositoryClient(ctx context.Context, server api.GitalyServer, addMetadataFuncs ...MetadataFunc) (context.Context, *RepositoryClient, error) {
|
2020-12-02 15:09:37 +00:00
|
|
|
conn, err := getOrCreateConnection(server)
|
|
|
|
if err != nil {
|
|
|
|
return nil, nil, err
|
|
|
|
}
|
|
|
|
grpcClient := gitalypb.NewRepositoryServiceClient(conn)
|
2022-05-18 09:07:58 +00:00
|
|
|
return withOutgoingMetadata(ctx, addMetadataFuncs...), &RepositoryClient{grpcClient}, nil
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
// NewNamespaceClient is only used by the Gitaly integration tests at present
|
2022-05-18 09:07:58 +00:00
|
|
|
func NewNamespaceClient(ctx context.Context, server api.GitalyServer, addMetadataFuncs ...MetadataFunc) (context.Context, *NamespaceClient, error) {
|
2020-12-02 15:09:37 +00:00
|
|
|
conn, err := getOrCreateConnection(server)
|
|
|
|
if err != nil {
|
|
|
|
return nil, nil, err
|
|
|
|
}
|
|
|
|
grpcClient := gitalypb.NewNamespaceServiceClient(conn)
|
2022-05-18 09:07:58 +00:00
|
|
|
return withOutgoingMetadata(ctx, addMetadataFuncs...), &NamespaceClient{grpcClient}, nil
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
func NewDiffClient(ctx context.Context, server api.GitalyServer, addMetadataFuncs ...MetadataFunc) (context.Context, *DiffClient, error) {
|
2020-12-02 15:09:37 +00:00
|
|
|
conn, err := getOrCreateConnection(server)
|
|
|
|
if err != nil {
|
|
|
|
return nil, nil, err
|
|
|
|
}
|
|
|
|
grpcClient := gitalypb.NewDiffServiceClient(conn)
|
2022-05-18 09:07:58 +00:00
|
|
|
return withOutgoingMetadata(ctx, addMetadataFuncs...), &DiffClient{grpcClient}, nil
|
2020-12-02 15:09:37 +00:00
|
|
|
}
|
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
func getOrCreateConnection(server api.GitalyServer) (*grpc.ClientConn, error) {
|
|
|
|
key := getCacheKey(server)
|
2020-12-02 15:09:37 +00:00
|
|
|
|
|
|
|
cache.RLock()
|
|
|
|
conn := cache.connections[key]
|
|
|
|
cache.RUnlock()
|
|
|
|
|
|
|
|
if conn != nil {
|
|
|
|
return conn, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
cache.Lock()
|
|
|
|
defer cache.Unlock()
|
|
|
|
|
|
|
|
if conn := cache.connections[key]; conn != nil {
|
|
|
|
return conn, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
conn, err := newConnection(server)
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
cache.connections[key] = conn
|
|
|
|
|
|
|
|
return conn, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func CloseConnections() {
|
|
|
|
cache.Lock()
|
|
|
|
defer cache.Unlock()
|
|
|
|
|
|
|
|
for _, conn := range cache.connections {
|
|
|
|
conn.Close()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-05-18 09:07:58 +00:00
|
|
|
func newConnection(server api.GitalyServer) (*grpc.ClientConn, error) {
|
2020-12-02 15:09:37 +00:00
|
|
|
connOpts := append(gitalyclient.DefaultDialOpts,
|
|
|
|
grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(server.Token)),
|
|
|
|
grpc.WithStreamInterceptor(
|
|
|
|
grpc_middleware.ChainStreamClient(
|
|
|
|
grpctracing.StreamClientTracingInterceptor(),
|
|
|
|
grpc_prometheus.StreamClientInterceptor,
|
|
|
|
grpccorrelation.StreamClientCorrelationInterceptor(
|
|
|
|
grpccorrelation.WithClientName("gitlab-workhorse"),
|
|
|
|
),
|
|
|
|
),
|
|
|
|
),
|
|
|
|
|
|
|
|
grpc.WithUnaryInterceptor(
|
|
|
|
grpc_middleware.ChainUnaryClient(
|
|
|
|
grpctracing.UnaryClientTracingInterceptor(),
|
|
|
|
grpc_prometheus.UnaryClientInterceptor,
|
|
|
|
grpccorrelation.UnaryClientCorrelationInterceptor(
|
|
|
|
grpccorrelation.WithClientName("gitlab-workhorse"),
|
|
|
|
),
|
|
|
|
),
|
|
|
|
),
|
|
|
|
)
|
|
|
|
|
2022-04-28 12:09:06 +00:00
|
|
|
conn, connErr := gitalyclient.DialSidechannel(context.Background(), server.Address, sidechannelRegistry, connOpts) // lint:allow context.Background
|
2020-12-02 15:09:37 +00:00
|
|
|
|
|
|
|
label := "ok"
|
|
|
|
if connErr != nil {
|
|
|
|
label = "fail"
|
|
|
|
}
|
|
|
|
connectionsTotal.WithLabelValues(label).Inc()
|
|
|
|
|
|
|
|
return conn, connErr
|
|
|
|
}
|
|
|
|
|
|
|
|
func UnmarshalJSON(s string, msg proto.Message) error {
|
|
|
|
return jsonUnMarshaler.Unmarshal(strings.NewReader(s), msg)
|
|
|
|
}
|