[Disk Manager] Support creation image from NRD disks

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

@@ -865,6 +865,44 @@ func (c *client) DeleteCheckpointData(
 	return err
 }
 
+func (c *client) EnsureCheckpointReady(
+	ctx context.Context,
+	diskID string,
+	checkpointID string,
+) (err error) {
+
+	defer c.metrics.StatRequest("EnsureCheckpointReady")(&err)
+
+	status, err := c.GetCheckpointStatus(ctx, diskID, checkpointID)
+
+	if err != nil {
+		return err
+	}
+
+	logging.Info(
+		ctx,
+		"Current CheckpointStatus: %v",
+		status,
+	)
+
+	switch status {
+	case CheckpointStatusNotReady:
+		return errors.NewInterruptExecutionError()
+
+	case CheckpointStatusError:
+		_ = c.DeleteCheckpoint(ctx, diskID, checkpointID)
+		return errors.NewRetriableErrorf(
+			"Creating checkpoint for disk with id %v ended with an error.",
+			diskID,
+		)
+
+	case CheckpointStatusReady:
+		// Nothing to do.
+	}
+
+	return nil
+}
+
 func (c *client) Resize(
 	ctx context.Context,
 	checkpoint func() error,

cloud/disk_manager/internal/pkg/clients/nbs/interface.go

@@ -161,6 +161,12 @@ type Client interface {
 		checkpointID string,
 	) error
 
+	EnsureCheckpointReady(
+		ctx context.Context,
+		diskID string,
+		checkpointID string,
+	) error
+
 	Resize(
 		ctx context.Context,
 		saveState func() error,

cloud/disk_manager/internal/pkg/clients/nbs/mocks/client_mock.go

@@ -106,6 +106,16 @@ func (c *ClientMock) DeleteCheckpointData(
 	return args.Error(0)
 }
 
+func (c *ClientMock) EnsureCheckpointReady(
+	ctx context.Context,
+	diskID string,
+	checkpointID string,
+) error {
+
+	args := c.Called(ctx, diskID, checkpointID)
+	return args.Error(0)
+}
+
 func (c *ClientMock) Resize(
 	ctx context.Context,
 	checkpoint func() error,

cloud/disk_manager/internal/pkg/facade/image_service_test/image_service_test.go

@@ -176,6 +176,91 @@ func checkUnencryptedImage(
 	require.ErrorContains(t, err, "KeyPath should contain path to encryption key")
 }
 
+func testImageServiceCreateImageFromDiskKind(
+	t *testing.T,
+	diskKind disk_manager.DiskKind,
+	diskSize uint64,
+) {
+
+	ctx := testcommon.NewContext()
+
+	client, err := testcommon.NewClient(ctx)
+	require.NoError(t, err)
+	defer client.Close()
+
+	diskID := t.Name()
+
+	reqCtx := testcommon.GetRequestContext(t, ctx)
+	operation, err := client.CreateDisk(reqCtx, &disk_manager.CreateDiskRequest{
+		Src: &disk_manager.CreateDiskRequest_SrcEmpty{
+			SrcEmpty: &empty.Empty{},
+		},
+		Size: int64(diskSize),
+		Kind: diskKind,
+		DiskId: &disk_manager.DiskId{
+			ZoneId: "zone-a",
+			DiskId: diskID,
+		},
+	})
+	require.NoError(t, err)
+	require.NotEmpty(t, operation)
+	err = internal_client.WaitOperation(ctx, client, operation.Id)
+	require.NoError(t, err)
+
+	nbsClient := testcommon.NewNbsClient(t, ctx, "zone-a")
+	diskContentInfo, err := nbsClient.FillDisk(ctx, diskID, diskSize)
+	require.NoError(t, err)
+
+	imageID := t.Name()
+
+	reqCtx = testcommon.GetRequestContext(t, ctx)
+	operation, err = client.CreateImage(reqCtx, &disk_manager.CreateImageRequest{
+		Src: &disk_manager.CreateImageRequest_SrcDiskId{
+			SrcDiskId: &disk_manager.DiskId{
+				ZoneId: "zone-a",
+				DiskId: diskID,
+			},
+		},
+		DstImageId: imageID,
+		FolderId:   "folder",
+		Pooled:     true,
+	})
+	require.NoError(t, err)
+	require.NotEmpty(t, operation)
+
+	response := disk_manager.CreateImageResponse{}
+	err = internal_client.WaitResponse(ctx, client, operation.Id, &response)
+	require.NoError(t, err)
+	require.Equal(t, int64(diskSize), response.Size)
+
+	meta := disk_manager.CreateImageMetadata{}
+	err = internal_client.GetOperationMetadata(ctx, client, operation.Id, &meta)
+	require.NoError(t, err)
+	require.Equal(t, float64(1), meta.Progress)
+
+	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID)
+
+	checkUnencryptedImage(
+		t,
+		client,
+		ctx,
+		imageID,
+		int64(diskSize),
+		diskContentInfo.Crc32,
+	)
+
+	reqCtx = testcommon.GetRequestContext(t, ctx)
+	operation, err = client.DeleteImage(reqCtx, &disk_manager.DeleteImageRequest{
+		ImageId: imageID,
+	})
+	require.NoError(t, err)
+	require.NotEmpty(t, operation)
+	err = internal_client.WaitOperation(ctx, client, operation.Id)
+	require.NoError(t, err)
+
+	testcommon.CheckConsistency(t, ctx)
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 func TestImageServiceCreateImageFromImage(t *testing.T) {
@@ -906,84 +991,19 @@ func TestImageServiceCancelCreateImageFromSnapshot(t *testing.T) {
 }
 
 func TestImageServiceCreateImageFromDisk(t *testing.T) {
-	ctx := testcommon.NewContext()
-
-	client, err := testcommon.NewClient(ctx)
-	require.NoError(t, err)
-	defer client.Close()
-
-	diskID := t.Name()
-	diskSize := uint64(4194304)
-
-	reqCtx := testcommon.GetRequestContext(t, ctx)
-	operation, err := client.CreateDisk(reqCtx, &disk_manager.CreateDiskRequest{
-		Src: &disk_manager.CreateDiskRequest_SrcEmpty{
-			SrcEmpty: &empty.Empty{},
-		},
-		Size: int64(diskSize),
-		Kind: disk_manager.DiskKind_DISK_KIND_SSD,
-		DiskId: &disk_manager.DiskId{
-			ZoneId: "zone-a",
-			DiskId: diskID,
-		},
-	})
-	require.NoError(t, err)
-	require.NotEmpty(t, operation)
-	err = internal_client.WaitOperation(ctx, client, operation.Id)
-	require.NoError(t, err)
-
-	nbsClient := testcommon.NewNbsClient(t, ctx, "zone-a")
-	diskContentInfo, err := nbsClient.FillDisk(ctx, diskID, diskSize)
-	require.NoError(t, err)
-
-	imageID := t.Name()
-
-	reqCtx = testcommon.GetRequestContext(t, ctx)
-	operation, err = client.CreateImage(reqCtx, &disk_manager.CreateImageRequest{
-		Src: &disk_manager.CreateImageRequest_SrcDiskId{
-			SrcDiskId: &disk_manager.DiskId{
-				ZoneId: "zone-a",
-				DiskId: diskID,
-			},
-		},
-		DstImageId: imageID,
-		FolderId:   "folder",
-		Pooled:     true,
-	})
-	require.NoError(t, err)
-	require.NotEmpty(t, operation)
-
-	response := disk_manager.CreateImageResponse{}
-	err = internal_client.WaitResponse(ctx, client, operation.Id, &response)
-	require.NoError(t, err)
-	require.Equal(t, int64(diskSize), response.Size)
-
-	meta := disk_manager.CreateImageMetadata{}
-	err = internal_client.GetOperationMetadata(ctx, client, operation.Id, &meta)
-	require.NoError(t, err)
-	require.Equal(t, float64(1), meta.Progress)
-
-	testcommon.RequireCheckpointsAreEmpty(t, ctx, diskID)
-
-	checkUnencryptedImage(
+	testImageServiceCreateImageFromDiskKind(
 		t,
-		client,
-		ctx,
-		imageID,
-		int64(diskSize),
-		diskContentInfo.Crc32,
+		disk_manager.DiskKind_DISK_KIND_SSD,
+		uint64(4194304),
 	)
+}
 
-	reqCtx = testcommon.GetRequestContext(t, ctx)
-	operation, err = client.DeleteImage(reqCtx, &disk_manager.DeleteImageRequest{
-		ImageId: imageID,
-	})
-	require.NoError(t, err)
-	require.NotEmpty(t, operation)
-	err = internal_client.WaitOperation(ctx, client, operation.Id)
-	require.NoError(t, err)
-
-	testcommon.CheckConsistency(t, ctx)
+func TestImageServiceCreateImageFromNonReplicatedDisk(t *testing.T) {
+	testImageServiceCreateImageFromDiskKind(
+		t,
+		disk_manager.DiskKind_DISK_KIND_SSD_NONREPLICATED,
+		uint64(1073741824),
+	)
 }
 
 func TestImageServiceCancelCreateImageFromDisk(t *testing.T) {

cloud/disk_manager/internal/pkg/services/images/create_image_from_disk_task.go

@@ -93,6 +93,11 @@ func (t *createImageFromDiskTask) run(
 		return err
 	}
 
+	err = nbsClient.EnsureCheckpointReady(ctx, disk.DiskId, checkpointID)
+	if err != nil {
+		return err
+	}
+
 	taskID, err := t.scheduler.ScheduleZonalTask(
 		headers.SetIncomingIdempotencyKey(ctx, selfTaskID+"_run"),
 		"dataplane.CreateSnapshotFromDisk",

cloud/disk_manager/internal/pkg/services/images/create_image_from_disk_task_test.go

@@ -0,0 +1,86 @@
+package images
+
+import (
+	"context"
+	"testing"
+
+	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/clients/nbs"
+	nbs_mocks "github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/clients/nbs/mocks"
+	performance_config "github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/performance/config"
+	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/resources"
+	storage_mocks "github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/resources/mocks"
+	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/services/images/protos"
+	"github.com/ydb-platform/nbs/cloud/disk_manager/internal/pkg/types"
+	"github.com/ydb-platform/nbs/cloud/tasks/errors"
+	tasks_mocks "github.com/ydb-platform/nbs/cloud/tasks/mocks"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+////////////////////////////////////////////////////////////////////////////////
+
+func newExecutionContextMock() *tasks_mocks.ExecutionContextMock {
+	execCtx := tasks_mocks.NewExecutionContextMock()
+	execCtx.On("GetTaskID").Return("toplevel_task_id")
+	return execCtx
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+func TestCreateImageFromDiskTaskFailure(t *testing.T) {
+	ctx := context.Background()
+	scheduler := tasks_mocks.NewSchedulerMock()
+	storage := storage_mocks.NewStorageMock()
+	nbsFactory := nbs_mocks.NewFactoryMock()
+	nbsClient := nbs_mocks.NewClientMock()
+	execCtx := newExecutionContextMock()
+
+	diskID := t.Name()
+	zoneID := "zone"
+	imageID := "image"
+	checkpointEnsuringError := errors.NewRetriableErrorf(
+		"Creating checkpoint for disk with id %v ended with an error.",
+		diskID,
+	)
+
+	diskParams := nbs.DiskParams{}
+	request := &protos.CreateImageFromDiskRequest{
+		SrcDisk: &types.Disk{
+			ZoneId: zoneID,
+			DiskId: diskID,
+		},
+		DstImageId: imageID,
+	}
+
+	task := &createImageFromDiskTask{
+		performanceConfig: &performance_config.PerformanceConfig{},
+		storage:           storage,
+		scheduler:         scheduler,
+		nbsFactory:        nbsFactory,
+		request:           request,
+		state:             &protos.CreateImageFromDiskTaskState{},
+	}
+
+	nbsFactory.On("GetClient", ctx, zoneID).Return(nbsClient, nil)
+	nbsClient.On("Describe", ctx, diskID).Return(diskParams, nil)
+
+	storage.On("CreateImage", ctx, mock.Anything).Return(
+		resources.ImageMeta{
+			Ready: false,
+		},
+		nil,
+	)
+
+	nbsClient.On("CreateCheckpoint", ctx, nbs.CheckpointParams{
+		DiskID:       diskID,
+		CheckpointID: imageID,
+	}).Return(nil)
+
+	nbsClient.On("EnsureCheckpointReady", ctx, diskID, imageID).Return(checkpointEnsuringError)
+
+	err := task.Run(ctx, execCtx)
+
+	require.Equal(t, err, checkpointEnsuringError)
+	mock.AssertExpectationsForObjects(t, storage, scheduler, nbsFactory, nbsClient, execCtx)
+}

cloud/disk_manager/internal/pkg/services/images/tests/ya.make

@@ -1,3 +1,8 @@
 GO_TEST_FOR(cloud/disk_manager/internal/pkg/services/images)
 
+REQUIREMENTS(
+    cpu:4
+    ram:16
+)
+
 END()

cloud/disk_manager/internal/pkg/services/images/ya.make

@@ -14,6 +14,7 @@ SRCS(
 )
 
 GO_TEST_SRCS(
+    create_image_from_disk_task_test.go
     delete_image_task_test.go
 )