From fd5c39a418bb2826763d5b4c3145762444264afb Mon Sep 17 00:00:00 2001
From: Zexi Li <zexi.li@icloud.com>
Date: Tue, 12 Apr 2022 13:53:26 +0800
Subject: [PATCH] fix(region): multi disk created after instance snapshot

---
 pkg/compute/models/instance_snapshot_joint.go | 10 +++++
 pkg/compute/models/instance_snapshots.go      | 37 +++++++++++++++++++
 pkg/compute/regiondrivers/kvm.go              | 19 ++++++----
 3 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/pkg/compute/models/instance_snapshot_joint.go b/pkg/compute/models/instance_snapshot_joint.go
index a940d5d7b8..a60b573fa7 100644
--- a/pkg/compute/models/instance_snapshot_joint.go
+++ b/pkg/compute/models/instance_snapshot_joint.go
@@ -17,6 +17,8 @@ package models
 import (
 	"context"
 
+	"yunion.io/x/pkg/errors"
+
 	"yunion.io/x/onecloud/pkg/cloudcommon/db"
 	"yunion.io/x/onecloud/pkg/mcclient"
 )
@@ -82,3 +84,11 @@ func (manager *SInstanceSnapshotJointManager) IsSubSnapshot(snapshotId string) (
 func (self *SInstanceSnapshotJoint) Detach(ctx context.Context, userCred mcclient.TokenCredential) error {
 	return db.DetachJoint(ctx, userCred, self)
 }
+
+func (self *SInstanceSnapshotJoint) GetSnapshotDisk() (*SDisk, error) {
+	sp, err := SnapshotManager.FetchById(self.SnapshotId)
+	if err != nil {
+		return nil, errors.Wrapf(err, "Get snapshot by %q", self.SnapshotId)
+	}
+	return sp.(*SSnapshot).GetDisk()
+}
diff --git a/pkg/compute/models/instance_snapshots.go b/pkg/compute/models/instance_snapshots.go
index ee3c5c76de..a41ddf1a9e 100644
--- a/pkg/compute/models/instance_snapshots.go
+++ b/pkg/compute/models/instance_snapshots.go
@@ -682,3 +682,40 @@ func (ism *SInstanceSnapshotManager) InitializeData() error {
 	}
 	return nil
 }
+
+func (isp *SInstanceSnapshot) GetInstanceSnapshotJointsByOrder(guest *SGuest) ([]*SInstanceSnapshotJoint, error) {
+	disks, err := guest.GetGuestDisks()
+	if err != nil {
+		return nil, errors.Wrap(err, "GetGuestDisks")
+	}
+	ss, err := isp.GetSnapshots()
+	if err != nil {
+		return nil, errors.Wrapf(err, "Get %s subsnapshots", isp.GetName())
+	}
+	jIsps := make([]*SInstanceSnapshotJoint, 0)
+	for idx, gd := range disks {
+		d := gd.GetDisk()
+		if d == nil {
+			return nil, errors.Wrapf(err, "Not get guestdisk %d related disk", idx)
+		}
+		if idx >= len(ss) {
+			break
+		}
+		jIsp, err := isp.GetInstanceSnapshotJointAt(idx)
+		if err != nil {
+			return nil, errors.Wrapf(err, "GetInstanceSnapshotJointAt %d", idx)
+		}
+		sd, err := ss[idx].GetDisk()
+		if err != nil {
+			return nil, errors.Wrapf(err, "Get snapshot %d disk", idx)
+		}
+		if ss[idx].GetId() != jIsp.SnapshotId {
+			return nil, errors.Wrapf(err, "InstanceSnapshotJoint %d snapshot_id %q != %q", idx, jIsp.SnapshotId, ss[idx].GetId())
+		}
+		if sd.GetId() != d.GetId() {
+			return nil, errors.Wrapf(err, "Disk Snapshot %d's disk id %q != current disk %q", idx, sd.GetId(), d.GetId())
+		}
+		jIsps = append(jIsps, jIsp)
+	}
+	return jIsps, nil
+}
diff --git a/pkg/compute/regiondrivers/kvm.go b/pkg/compute/regiondrivers/kvm.go
index 2b951f57d5..789024c424 100644
--- a/pkg/compute/regiondrivers/kvm.go
+++ b/pkg/compute/regiondrivers/kvm.go
@@ -1065,13 +1065,16 @@ func (self *SKVMRegionDriver) RequestDeleteInstanceBackup(ctx context.Context, i
 }
 
 func (self *SKVMRegionDriver) RequestResetToInstanceSnapshot(ctx context.Context, guest *models.SGuest, isp *models.SInstanceSnapshot, task taskman.ITask, params *jsonutils.JSONDict) error {
-	disks, _ := guest.GetGuestDisks()
+	jIsps, err := isp.GetInstanceSnapshotJointsByOrder(guest)
+	if err != nil {
+		return errors.Wrap(err, "GetInstanceSnapshotJointsByOrder")
+	}
 	diskIndexI64, err := params.Int("disk_index")
 	if err != nil {
 		return errors.Wrap(err, "get 'disk_index' from params")
 	}
 	diskIndex := int(diskIndexI64)
-	if diskIndex >= len(disks) {
+	if diskIndex >= len(jIsps) {
 		task.SetStage("OnInstanceSnapshotReset", nil)
 		withMem := jsonutils.QueryBoolean(params, "with_memory", false)
 		if isp.WithMemory && withMem {
@@ -1097,19 +1100,19 @@ func (self *SKVMRegionDriver) RequestResetToInstanceSnapshot(ctx context.Context
 		return nil
 	}
 
-	isj, err := isp.GetInstanceSnapshotJointAt(diskIndex)
-	if err != nil {
-		return err
-	}
+	isj := jIsps[diskIndex]
 
 	params = jsonutils.NewDict()
 	params.Set("disk_index", jsonutils.NewInt(int64(diskIndex)))
 	task.SetStage("OnKvmDiskReset", params)
 
-	disk := disks[diskIndex].GetDisk()
+	disk, err := isj.GetSnapshotDisk()
+	if err != nil {
+		return errors.Wrapf(err, "Get %d snapshot disk", diskIndex)
+	}
 	err = disk.StartResetDisk(ctx, task.GetUserCred(), isj.SnapshotId, false, guest, task.GetTaskId())
 	if err != nil {
-		return err
+		return errors.Wrap(err, "StartResetDisk")
 	}
 	return nil
 }