From 22b8fbafa48a583e901dddd851e833c123e7ea75 Mon Sep 17 00:00:00 2001
From: Moath Qasim <moad.qassem@gmail.com>
Date: Tue, 15 Sep 2020 09:25:41 +0200
Subject: [PATCH] referencing pre-allocated data volumes in vms

Signed-off-by: Moath Qasim <moad.qassem@gmail.com>
---
 cmd/machine-controller/main.go |  7 +++
 pkg/kubevirt/core/core.go      | 79 ++++++++++++++++++++++++----------
 pkg/kubevirt/core/core_test.go | 10 +++++
 3 files changed, 73 insertions(+), 23 deletions(-)

diff --git a/cmd/machine-controller/main.go b/cmd/machine-controller/main.go
index 410061e..c9c05ad 100644
--- a/cmd/machine-controller/main.go
+++ b/cmd/machine-controller/main.go
@@ -26,8 +26,10 @@ import (
 	_ "github.com/gardener/machine-controller-manager/pkg/util/reflector/prometheus" // for reflector metric registration
 	_ "github.com/gardener/machine-controller-manager/pkg/util/workqueue/prometheus" // for workqueue metric registration
 	"github.com/spf13/pflag"
+	"k8s.io/client-go/kubernetes/scheme"
 	"k8s.io/component-base/cli/flag"
 	"k8s.io/component-base/logs"
+	cdi "kubevirt.io/containerized-data-importer/pkg/apis/core/v1alpha1"
 )
 
 func main() {
@@ -39,6 +41,11 @@ func main() {
 	logs.InitLogs()
 	defer logs.FlushLogs()
 
+	if err := cdi.AddToScheme(scheme.Scheme); err != nil {
+		fmt.Fprintf(os.Stderr, " %v\n", err)
+		os.Exit(1)
+	}
+
 	plugin := kubevirt.NewKubevirtPlugin()
 
 	if err := app.Run(s, plugin); err != nil {
diff --git a/pkg/kubevirt/core/core.go b/pkg/kubevirt/core/core.go
index 0488ce7..423ff40 100644
--- a/pkg/kubevirt/core/core.go
+++ b/pkg/kubevirt/core/core.go
@@ -40,8 +40,11 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
-// ProviderName specifies the machine controller for kubevirt cloud provider
-const ProviderName = "kubevirt"
+const (
+	// ProviderName specifies the machine controller for kubevirt cloud provider
+	ProviderName      = "kubevirt"
+	machineClassLabel = "mcm.gardener.cloud/machineclass"
+)
 
 // ClientFactory creates a client from the kubeconfig saved in the "kubeconfig" field of the given secret.
 type ClientFactory interface {
@@ -134,6 +137,44 @@ func (p PluginSPIImpl) CreateMachine(ctx context.Context, machineName string, pr
 
 	vmLabels["kubevirt.io/vm"] = machineName
 
+	machineClassName := vmLabels[machineClassLabel]
+	dataVolumeName, err := p.getDataVolume(ctx, c, machineClassName, namespace)
+	if err != nil {
+		return "", err
+	}
+
+	dataVolumeTemplate := cdi.DataVolume{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      machineName,
+			Namespace: namespace,
+		},
+		Spec: cdi.DataVolumeSpec{
+			PVC: &corev1.PersistentVolumeClaimSpec{
+				StorageClassName: utilpointer.StringPtr(providerSpec.StorageClassName),
+				AccessModes: []corev1.PersistentVolumeAccessMode{
+					"ReadWriteOnce",
+				},
+				Resources: corev1.ResourceRequirements{
+					Requests: pvcRequest,
+				},
+			},
+			Source: cdi.DataVolumeSource{
+				HTTP: &cdi.DataVolumeSourceHTTP{
+					URL: providerSpec.SourceURL,
+				},
+			},
+		},
+	}
+
+	if dataVolumeName != "" {
+		dataVolumeTemplate.Spec.Source = cdi.DataVolumeSource{
+			PVC: &cdi.DataVolumeSourcePVC{
+				Name:      dataVolumeName,
+				Namespace: namespace,
+			},
+		}
+	}
+
 	virtualMachine := &kubevirtv1.VirtualMachine{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      machineName,
@@ -196,27 +237,7 @@ func (p PluginSPIImpl) CreateMachine(ctx context.Context, machineName string, pr
 				},
 			},
 			DataVolumeTemplates: []cdi.DataVolume{
-				{
-					ObjectMeta: metav1.ObjectMeta{
-						Name: machineName,
-					},
-					Spec: cdi.DataVolumeSpec{
-						PVC: &corev1.PersistentVolumeClaimSpec{
-							StorageClassName: utilpointer.StringPtr(providerSpec.StorageClassName),
-							AccessModes: []corev1.PersistentVolumeAccessMode{
-								"ReadWriteOnce",
-							},
-							Resources: corev1.ResourceRequirements{
-								Requests: pvcRequest,
-							},
-						},
-						Source: cdi.DataVolumeSource{
-							HTTP: &cdi.DataVolumeSourceHTTP{
-								URL: providerSpec.SourceURL,
-							},
-						},
-					},
-				},
+				dataVolumeTemplate,
 			},
 		},
 	}
@@ -341,3 +362,15 @@ func (p PluginSPIImpl) machineProviderID(ctx context.Context, c client.Client, m
 
 	return encodeProviderID(string(virtualMachine.UID)), nil
 }
+
+func (p PluginSPIImpl) getDataVolume(ctx context.Context, c client.Client, dataVolumeName, namespace string) (string, error) {
+	dataVolume := &cdi.DataVolume{}
+	if err := c.Get(ctx, types.NamespacedName{Namespace: namespace, Name: dataVolumeName}, dataVolume); err != nil {
+		if kerrors.IsNotFound(err) {
+			return "", nil
+		}
+		return "", fmt.Errorf("failed to get DataVolume: %v", err)
+	}
+
+	return dataVolume.Name, nil
+}
diff --git a/pkg/kubevirt/core/core_test.go b/pkg/kubevirt/core/core_test.go
index 965973c..d0897dc 100644
--- a/pkg/kubevirt/core/core_test.go
+++ b/pkg/kubevirt/core/core_test.go
@@ -16,12 +16,15 @@ package core
 
 import (
 	"context"
+	"os"
 	"testing"
 
 	api "github.com/gardener/machine-controller-manager-provider-kubevirt/pkg/kubevirt/apis"
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes/scheme"
+	"k8s.io/klog"
+	cdi "kubevirt.io/containerized-data-importer/pkg/apis/core/v1alpha1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 )
@@ -38,6 +41,13 @@ var (
 	namespace   = "default"
 )
 
+func init() {
+	if err := cdi.AddToScheme(scheme.Scheme); err != nil {
+		klog.Errorf("could not execute tests: %v", err)
+		os.Exit(1)
+	}
+}
+
 func TestPluginSPIImpl_CreateMachine(t *testing.T) {
 	fakeClient := fake.NewFakeClientWithScheme(scheme.Scheme)
 	t.Run("CreateMachine", func(t *testing.T) {