Ota Upload on multiple devices

Author: polldo

cli/ota/upload.go

@@ -28,9 +28,11 @@ import (
 )
 
 var uploadFlags struct {
-	deviceID string
-	file     string
-	deferred bool
+	deviceIDs []string
+	tags      map[string]string
+	file      string
+	deferred  bool
+	fqbn      string
 }
 
 func initUploadCommand() *cobra.Command {
@@ -41,23 +43,33 @@ func initUploadCommand() *cobra.Command {
 		Run:   runUploadCommand,
 	}
 
-	uploadCommand.Flags().StringVarP(&uploadFlags.deviceID, "device-id", "d", "", "Device ID")
+	uploadCommand.Flags().StringSliceVarP(&uploadFlags.deviceIDs, "device-ids", "d", []string{},
+		"Comma-separated list of device IDs to update")
+	uploadCommand.Flags().StringToStringVar(&uploadFlags.tags, "tags", nil,
+		"Comma-separated list of tags with format <key>=<value>.\n"+
+			"Perform and OTA upload on all devices that match the provided tags.\n"+
+			"Mutually exclusive with `--device-id`.",
+	)
 	uploadCommand.Flags().StringVarP(&uploadFlags.file, "file", "", "", "Binary file (.bin) to be uploaded")
 	uploadCommand.Flags().BoolVar(&uploadFlags.deferred, "deferred", false, "Perform a deferred OTA. It can take up to 1 week.")
+	uploadCommand.Flags().StringVarP(&uploadFlags.fqbn, "fqbn", "b", "", "FQBN of the devices to update")
 
-	uploadCommand.MarkFlagRequired("device-id")
 	uploadCommand.MarkFlagRequired("file")
+	uploadCommand.MarkFlagRequired("fqbn")
 	return uploadCommand
 }
 
 func runUploadCommand(cmd *cobra.Command, args []string) {
-	logrus.Infof("Uploading binary %s to device %s", uploadFlags.file, uploadFlags.deviceID)
+	logrus.Infof("Uploading binary %s", uploadFlags.file)
 
 	params := &ota.UploadParams{
-		DeviceID: uploadFlags.deviceID,
-		File:     uploadFlags.file,
-		Deferred: uploadFlags.deferred,
+		DeviceIDs: uploadFlags.deviceIDs,
+		Tags:      uploadFlags.tags,
+		File:      uploadFlags.file,
+		Deferred:  uploadFlags.deferred,
+		FQBN:      uploadFlags.fqbn,
 	}
+
 	err := ota.Upload(params)
 	if err != nil {
 		feedback.Errorf("Error during ota upload: %v", err)

command/ota/upload.go

@@ -18,10 +18,12 @@
 package ota
 
 import (
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/arduino/arduino-cloud-cli/internal/config"
 	"github.com/arduino/arduino-cloud-cli/internal/iot"
@@ -32,19 +34,29 @@ const (
 	otaExpirationMins = 10
 	// deferred ota can take up to 1 week (equal to 10080 minutes)
 	otaDeferredExpirationMins = 10080
+
+	numConcurrentUploads = 10
 )
 
 // UploadParams contains the parameters needed to
 // perform an OTA upload.
 type UploadParams struct {
-	DeviceID string
-	File     string
-	Deferred bool
+	DeviceIDs []string
+	Tags      map[string]string
+	File      string
+	Deferred  bool
+	FQBN      string
 }
 
 // Upload command is used to upload a firmware OTA,
 // on a device of Arduino IoT Cloud.
 func Upload(params *UploadParams) error {
+	if params.DeviceIDs == nil && params.Tags == nil {
+		return errors.New("provide either DeviceID or Tags")
+	} else if params.DeviceIDs != nil && params.Tags != nil {
+		return errors.New("cannot use both DeviceID and Tags. only one of them should be not nil")
+	}
+
 	conf, err := config.Retrieve()
 	if err != nil {
 		return err
@@ -54,10 +66,14 @@ func Upload(params *UploadParams) error {
 		return err
 	}
 
-	dev, err := iotClient.DeviceShow(params.DeviceID)
+	d, err := idsGivenTags(iotClient, params.Tags)
 	if err != nil {
 		return err
 	}
+	devs := append(params.DeviceIDs, d...)
+	if len(devs) == 0 {
+		return errors.New("no device found")
+	}
 
 	otaDir, err := ioutil.TempDir("", "")
 	if err != nil {
@@ -66,7 +82,7 @@ func Upload(params *UploadParams) error {
 	otaFile := filepath.Join(otaDir, "temp.ota")
 	defer os.RemoveAll(otaDir)
 
-	err = Generate(params.File, otaFile, dev.Fqbn)
+	err = Generate(params.File, otaFile, params.FQBN)
 	if err != nil {
 		return fmt.Errorf("%s: %w", "cannot generate .ota file", err)
 	}
@@ -81,10 +97,56 @@ func Upload(params *UploadParams) error {
 		expiration = otaDeferredExpirationMins
 	}
 
-	err = iotClient.DeviceOTA(params.DeviceID, file, expiration)
+	return run(iotClient, devs, file, expiration)
+}
+
+func idsGivenTags(iotClient iot.Client, tags map[string]string) ([]string, error) {
+	if tags == nil {
+		return nil, nil
+	}
+	devs, err := iotClient.DeviceList(tags)
 	if err != nil {
-		return err
+		return nil, fmt.Errorf("%s: %w", "cannot retrieve devices from cloud", err)
+	}
+	devices := make([]string, 0, len(devs))
+	for _, d := range devs {
+		devices = append(devices, d.Id)
+	}
+	return devices, nil
+}
+
+func run(iotClient iot.Client, ids []string, file *os.File, expiration int) error {
+	idsToProcess := make(chan string, 2000)
+	idsFailed := make(chan string, 2000)
+	for _, id := range ids {
+		idsToProcess <- id
+	}
+	close(idsToProcess)
+
+	for i := 0; i < numConcurrentUploads; i++ {
+		go func() {
+			for id := range idsToProcess {
+				err := iotClient.DeviceOTA(id, file, expiration)
+				fail := ""
+				if err != nil {
+					fail = id
+				}
+				idsFailed <- fail
+			}
+		}()
+	}
+
+	failMsg := ""
+	for range ids {
+		i := <-idsFailed
+		if i != "" {
+			failMsg = strings.Join([]string{i, failMsg}, ",")
+		}
 	}
 
+	if failMsg != "" {
+		failMsg = strings.TrimRight(failMsg, ",")
+		return fmt.Errorf("failed to update these devices: %s", failMsg)
+	}
 	return nil
 }

command/ota/upload_test.go

@@ -0,0 +1,36 @@
+package ota
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/arduino/arduino-cloud-cli/internal/iot/mocks"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestRun(t *testing.T) {
+	mockClient := &mocks.Client{}
+	mockDeviceOTA := func(id string, file *os.File, expireMins int) error {
+		time.Sleep(3 * time.Second)
+		if strings.Split(id, "-")[0] == "fail" {
+			return errors.New("err")
+		}
+		return nil
+	}
+	mockClient.On("DeviceOTA", mock.Anything, mock.Anything, mock.Anything).Return(mockDeviceOTA, nil)
+
+	err := run(mockClient, []string{"dont-fail", "fail-1", "dont-fail", "fail-2"}, nil, 0)
+	if err == nil {
+		t.Error("should return error")
+	}
+	fmt.Println(err.Error())
+	failed := strings.Split(err.Error(), ",")
+	if len(failed) != 2 {
+		fmt.Println(len(failed), failed)
+		t.Error("two updates should have failed")
+	}
+}