TMP: azure: Save screenshot on failure

Author: Jeremi Piotrowski

platform/api/azure/instance.go

@@ -30,12 +30,20 @@ import (
 	"github.com/flatcar/mantle/util"
 )
 
+type MachineState int
+
+const (
+	READY MachineState = iota
+	PROVISIONING
+)
+
 type Machine struct {
 	ID               string
 	PublicIPAddress  string
 	PrivateIPAddress string
 	InterfaceName    string
 	PublicIPName     string
+	State            MachineState
 }
 
 func (a *API) getAvset() string {
@@ -239,7 +247,7 @@ func (a *API) CreateInstance(name, sshkey, resourceGroup, storageAccount string,
 	defer cancel()
 	_, err = poller.PollUntilDone(ctx, nil)
 	if err != nil {
-		return &Machine{ID: name}, fmt.Errorf("PollUntilDone: %w", err)
+		return &Machine{ID: name, State: PROVISIONING}, fmt.Errorf("PollUntilDone: %w", err)
 	}
 	plog.Infof("Instance %s created", name)
 
@@ -300,6 +308,38 @@ func (a *API) TerminateInstance(machine *Machine, resourceGroup string) error {
 	return err
 }
 
+func (a *API) GetScreenshot(name, resourceGroup string, output io.Writer) error {
+	vmResourceGroup := a.getVMRG(resourceGroup)
+	param := &armcompute.VirtualMachinesClientRetrieveBootDiagnosticsDataOptions{
+		SasURIExpirationTimeInMinutes: to.Ptr[int32](5),
+	}
+	resp, err := a.compClient.RetrieveBootDiagnosticsData(context.TODO(), vmResourceGroup, name, param)
+	if err != nil {
+		return fmt.Errorf("could not get VM: %v", err)
+	}
+	if resp.ConsoleScreenshotBlobURI == nil {
+		return fmt.Errorf("console screenshot URI is nil")
+	}
+
+	var data io.ReadCloser
+	err = util.Retry(6, 10*time.Second, func() error {
+		reply, err := http.Get(*resp.ConsoleScreenshotBlobURI)
+		if err != nil {
+			return fmt.Errorf("could not GET console screenshot: %v", err)
+		}
+		data = reply.Body
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	written, err := io.Copy(output, data)
+	if err == nil {
+		plog.Debugf("wrote %d bytes to screenshot", written)
+	}
+	return err
+}
+
 func (a *API) GetConsoleOutput(name, resourceGroup, storageAccount string) ([]byte, error) {
 	vmResourceGroup := a.getVMRG(resourceGroup)
 	param := &armcompute.VirtualMachinesClientRetrieveBootDiagnosticsDataOptions{

platform/machine/azure/machine.go

@@ -126,6 +126,16 @@ func (am *machine) saveConsole() error {
 		return fmt.Errorf("failed writing console to file: %v", err)
 	}
 
+	if am.mach.State == azure.PROVISIONING {
+		path := filepath.Join(am.dir, "screenshot.bmp")
+		f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0666)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		return am.cluster.flight.Api.GetScreenshot(am.ID(), am.ResourceGroup(), f)
+	}
+
 	return nil
 }