[Disk Manager] Add ssd-nbs2 disks created via YDB NbsService

cloud/disk_manager/api/disk.proto

@@ -15,6 +15,7 @@ enum DiskKind {
     DISK_KIND_SSD_MIRROR3 = 6;
     DISK_KIND_HDD_NONREPLICATED = 7;
     DISK_KIND_HDD_LOCAL = 8;
+    DISK_KIND_SSD_NBS2 = 9;
 }
 
 enum DiskState {

cloud/disk_manager/internal/pkg/cells/cells.go

@@ -370,7 +370,7 @@ func (s *cellSelector) getRecentAggregatedClusterCapacities(
 	zoneID string,
 ) ([]storage.ClusterCapacity, error) {
 
-	diskKinds := util.GetAllDiskKinds()
+	diskKinds := util.GetBlockstoreDiskKinds()
 
 	aggregated := make(map[string]*storage.ClusterCapacity)
 

cloud/disk_manager/internal/pkg/clients/nbs2/client.go

@@ -0,0 +1,326 @@
+package nbs2
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/golang/protobuf/proto"
+	nbs2_config "github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/clients/nbs2/config"
+	nbs2_protos "github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/clients/nbs2/protos"
+	"github.com/ydb-platform/nbs/cloud/tasks/errors"
+	"golang.org/x/exp/maps"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+////////////////////////////////////////////////////////////////////////////////
+
+// Ydb.StatusIds.StatusCode values used by Ydb.Nbs.V1.NbsService.
+const (
+	ydbStatusSuccess            int32 = 400000
+	ydbStatusBadRequest         int32 = 400010
+	ydbStatusUnauthorized       int32 = 400020
+	ydbStatusInternalError      int32 = 400030
+	ydbStatusAborted            int32 = 400040
+	ydbStatusUnavailable        int32 = 400050
+	ydbStatusOverloaded         int32 = 400060
+	ydbStatusSchemeError        int32 = 400070
+	ydbStatusGenericError       int32 = 400080
+	ydbStatusTimeout            int32 = 400090
+	ydbStatusBadSession         int32 = 400100
+	ydbStatusPreconditionFailed int32 = 400120
+	ydbStatusAlreadyExists      int32 = 400130
+	ydbStatusNotFound           int32 = 400140
+	ydbStatusSessionExpired     int32 = 400150
+	ydbStatusCancelled          int32 = 400160
+	ydbStatusUndetermined       int32 = 400170
+	ydbStatusUnsupported        int32 = 400180
+	ydbStatusSessionBusy        int32 = 400190
+	ydbStatusExternalError      int32 = 400200
+)
+
+type client struct {
+	zoneID  string
+	timeout time.Duration
+	dial    func(ctx context.Context) (*grpc.ClientConn, error)
+}
+
+func (c *client) ZoneID() string {
+	return c.zoneID
+}
+
+func (c *client) CreatePartition(
+	ctx context.Context,
+	params CreatePartitionParams,
+) (string, error) {
+
+	if len(params.DiskID) == 0 {
+		return "", errors.NewNonRetriableErrorf("disk id is required")
+	}
+	if len(params.StoragePoolName) == 0 {
+		return "", errors.NewNonRetriableErrorf(
+			"storage pool name is required for ssd-nbs2 disk %v",
+			params.DiskID,
+		)
+	}
+
+	req := &nbs2_protos.CreatePartitionRequest{
+		OperationParams: &nbs2_protos.OperationParams{
+			OperationMode: nbs2_protos.OperationParams_SYNC,
+		},
+		DiskId:          params.DiskID,
+		BlockSize:       params.BlockSize,
+		BlocksCount:     params.BlocksCount,
+		StoragePoolName: params.StoragePoolName,
+		StorageMedia:    nbs2_protos.StorageMediaKind_STORAGE_MEDIA_SSD,
+	}
+
+	resp := &nbs2_protos.CreatePartitionResponse{}
+	err := c.invoke(ctx, "CreatePartition", req, resp)
+	if err != nil {
+		return "", err
+	}
+
+	op, err := checkOperation(
+		resp.GetOperation(),
+		ydbStatusSuccess,
+		ydbStatusAlreadyExists,
+	)
+	if err != nil {
+		return "", err
+	}
+
+	result := &nbs2_protos.CreatePartitionResult{}
+	if op.GetResult() != nil && len(op.GetResult().GetValue()) > 0 {
+		err = unpackOperationResult(op, result)
+		if err != nil {
+			return "", err
+		}
+	}
+
+	return result.GetTabletId(), nil
+}
+
+func (c *client) DeletePartition(ctx context.Context, diskID string) error {
+	if len(diskID) == 0 {
+		return errors.NewNonRetriableErrorf("disk id is required")
+	}
+
+	req := &nbs2_protos.DeletePartitionRequest{
+		OperationParams: &nbs2_protos.OperationParams{
+			OperationMode: nbs2_protos.OperationParams_SYNC,
+		},
+		DiskId: diskID,
+	}
+
+	resp := &nbs2_protos.DeletePartitionResponse{}
+	err := c.invoke(ctx, "DeletePartition", req, resp)
+	if err != nil {
+		return err
+	}
+
+	_, err = checkOperation(
+		resp.GetOperation(),
+		ydbStatusSuccess,
+		ydbStatusNotFound,
+	)
+	return err
+}
+
+func (c *client) invoke(
+	ctx context.Context,
+	method string,
+	req proto.Message,
+	resp proto.Message,
+) error {
+
+	timeout := c.timeout
+	if timeout <= 0 {
+		timeout = 20 * time.Second
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	conn, err := c.dial(ctx)
+	if err != nil {
+		return errors.NewRetriableError(err)
+	}
+	defer conn.Close()
+
+	stub := nbs2_protos.NewNbsServiceClient(conn)
+	switch method {
+	case "CreatePartition":
+		out, err := stub.CreatePartition(ctx, req.(*nbs2_protos.CreatePartitionRequest))
+		if err != nil {
+			return errors.NewRetriableError(err)
+		}
+		proto.Merge(resp, out)
+	case "DeletePartition":
+		out, err := stub.DeletePartition(ctx, req.(*nbs2_protos.DeletePartitionRequest))
+		if err != nil {
+			return errors.NewRetriableError(err)
+		}
+		proto.Merge(resp, out)
+	default:
+		return errors.NewNonRetriableErrorf("unknown nbs method %v", method)
+	}
+
+	return nil
+}
+
+func checkOperation(
+	op *nbs2_protos.Operation,
+	okStatuses ...int32,
+) (*nbs2_protos.Operation, error) {
+
+	if op == nil {
+		return nil, errors.NewRetriableErrorf("empty operation in nbs response")
+	}
+	if !op.GetReady() {
+		return nil, errors.NewRetriableErrorf(
+			"nbs operation %v is not ready",
+			op.GetId(),
+		)
+	}
+	if len(okStatuses) == 0 {
+		okStatuses = []int32{ydbStatusSuccess}
+	}
+	for _, status := range okStatuses {
+		if op.GetStatus() == status {
+			return op, nil
+		}
+	}
+	return nil, operationStatusError(op)
+}
+
+func operationStatusError(op *nbs2_protos.Operation) error {
+	msg := fmt.Sprintf(
+		"nbs operation %v failed with status %v",
+		op.GetId(),
+		op.GetStatus(),
+	)
+	if isRetriableYdbStatus(op.GetStatus()) {
+		return errors.NewRetriableErrorf("%s", msg)
+	}
+	return errors.NewNonRetriableErrorf("%s", msg)
+}
+
+func isRetriableYdbStatus(status int32) bool {
+	switch status {
+	case ydbStatusInternalError,
+		ydbStatusAborted,
+		ydbStatusUnavailable,
+		ydbStatusOverloaded,
+		ydbStatusTimeout,
+		ydbStatusBadSession,
+		ydbStatusSessionExpired,
+		ydbStatusCancelled,
+		ydbStatusUndetermined,
+		ydbStatusSessionBusy:
+		return true
+	default:
+		return false
+	}
+}
+
+func unpackOperationResult(op *nbs2_protos.Operation, msg proto.Message) error {
+	if op.GetResult() == nil {
+		return errors.NewNonRetriableErrorf("nbs operation %v has empty result", op.GetId())
+	}
+
+	err := proto.Unmarshal(op.GetResult().GetValue(), msg)
+	if err != nil {
+		return errors.NewNonRetriableErrorf(
+			"failed to unpack nbs operation %v result: %w",
+			op.GetId(),
+			err,
+		)
+	}
+	return nil
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+type factory struct {
+	config  *nbs2_config.ClientConfig
+	timeout time.Duration
+}
+
+func (f *factory) GetClient(ctx context.Context, zoneID string) (Client, error) {
+	if f.config == nil {
+		return nil, errors.NewNonRetriableErrorf(
+			"nbs2 client is not configured, available zones: []",
+		)
+	}
+
+	zone, ok := f.config.GetZones()[zoneID]
+	if !ok {
+		return nil, errors.NewNonRetriableErrorf(
+			"unknown nbs2 zone %q, available zones: %q",
+			zoneID,
+			maps.Keys(f.config.GetZones()),
+		)
+	}
+	if len(zone.GetEndpoints()) == 0 {
+		return nil, errors.NewNonRetriableErrorf(
+			"no nbs2 endpoints for zone %q",
+			zoneID,
+		)
+	}
+
+	endpoint := normalizeEndpoint(zone.GetEndpoints()[0])
+	creds, err := f.transportCredentials()
+	if err != nil {
+		return nil, err
+	}
+
+	return &client{
+		zoneID:  zoneID,
+		timeout: f.timeout,
+		dial: func(ctx context.Context) (*grpc.ClientConn, error) {
+			return grpc.DialContext(ctx, endpoint, grpc.WithTransportCredentials(creds))
+		},
+	}, nil
+}
+
+func (f *factory) transportCredentials() (credentials.TransportCredentials, error) {
+	if f.config.GetInsecure() || len(f.config.GetRootCertsFile()) == 0 {
+		return insecure.NewCredentials(), nil
+	}
+
+	creds, err := credentials.NewClientTLSFromFile(f.config.GetRootCertsFile(), "")
+	if err != nil {
+		return nil, errors.NewNonRetriableErrorf(
+			"failed to load nbs2 root certs from %v: %w",
+			f.config.GetRootCertsFile(),
+			err,
+		)
+	}
+	return creds, nil
+}
+
+func normalizeEndpoint(endpoint string) string {
+	endpoint = strings.TrimPrefix(endpoint, "grpc://")
+	endpoint = strings.TrimPrefix(endpoint, "grpcs://")
+	return endpoint
+}
+
+func NewFactory(config *nbs2_config.ClientConfig) (Factory, error) {
+	timeout := 20 * time.Second
+	if config != nil && len(config.GetRequestTimeout()) > 0 {
+		parsed, err := time.ParseDuration(config.GetRequestTimeout())
+		if err != nil {
+			return nil, fmt.Errorf("invalid nbs2 request timeout: %w", err)
+		}
+		timeout = parsed
+	}
+
+	return &factory{
+		config:  config,
+		timeout: timeout,
+	}, nil
+}