RSDK-9060 Allow custom first run timeout configuration (#4511)

Co-authored-by: Benjamin Rewis <[email protected]>
viamrobotics · Nov 4, 2024 · 02e1b4a · 02e1b4a
1 parent b8d7747
commit 02e1b4a
Show file tree

Hide file tree

Showing 8 changed files with 217 additions and 121 deletions.
diff --git a/config/config_test.go b/config/config_test.go
@@ -17,6 +17,7 @@ import (
 	"github.com/lestrrat-go/jwx/jwk"
 	"github.com/pkg/errors"
 	"go.viam.com/test"
+	"go.viam.com/utils"
 	"go.viam.com/utils/jwks"
 	"go.viam.com/utils/pexec"
 	"go.viam.com/utils/rpc"
@@ -1220,22 +1221,70 @@ func TestConfigRobotRevision(t *testing.T) {
 	test.That(t, cfg.Revision, test.ShouldEqual, "rev1")
 }
 
-func TestConfigMarshalUnMarshal(t *testing.T) {
-	c := config.Config{
-		MaintenanceConfig: &config.MaintenanceConfig{
-			SensorName:            "SensorName",
-			MaintenanceAllowedKey: "Key",
-		},
+func TestConfigJSONMarshalRoundtrip(t *testing.T) {
+	type testcase struct {
+		name     string
+		c        config.Config
+		expected config.Config
 	}
-	expectedVal := config.Config{
-		MaintenanceConfig: &config.MaintenanceConfig{
-			SensorName:            "SensorName",
-			MaintenanceAllowedKey: "Key",
+
+	for _, tc := range []testcase{
+		{
+			name: "maintenance config",
+			c: config.Config{
+				MaintenanceConfig: &config.MaintenanceConfig{
+					SensorName:            "SensorName",
+					MaintenanceAllowedKey: "Key",
+				},
+			},
+			expected: config.Config{
+				MaintenanceConfig: &config.MaintenanceConfig{
+					SensorName:            "SensorName",
+					MaintenanceAllowedKey: "Key",
+				},
+			},
 		},
-	}
-	val, err := c.MarshalJSON()
-	test.That(t, err, test.ShouldBeNil)
+		{
+			name: "module",
+			c: config.Config{
+				Modules: []config.Module{
+					{
+						Name:            "ModuleName",
+						ExePath:         "ExecutablePath",
+						LogLevel:        "WARN",
+						Type:            config.ModuleTypeLocal,
+						ModuleID:        "ModuleID",
+						Environment:     map[string]string{"KEY": "VAL"},
+						FirstRunTimeout: utils.Duration(5 * time.Minute),
+						Status:          &config.AppValidationStatus{Error: "durrr"},
+					},
+				},
+			},
+			expected: config.Config{
+				Modules: []config.Module{
+					{
+						Name:            "ModuleName",
+						ExePath:         "ExecutablePath",
+						LogLevel:        "WARN",
+						Type:            config.ModuleTypeLocal,
+						ModuleID:        "ModuleID",
+						FirstRunTimeout: utils.Duration(5 * time.Minute),
+						Environment:     map[string]string{"KEY": "VAL"},
+						Status:          &config.AppValidationStatus{Error: "durrr"},
+					},
+				},
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			c := tc.c
+
+			data, err := c.MarshalJSON()
+			test.That(t, err, test.ShouldBeNil)
 
-	c.UnmarshalJSON(val)
-	test.That(t, c, test.ShouldResemble, expectedVal)
+			err = c.UnmarshalJSON(data)
+			test.That(t, err, test.ShouldBeNil)
+			test.That(t, c, test.ShouldResemble, tc.expected)
+		})
+	}
 }
diff --git a/config/module.go b/config/module.go
@@ -8,9 +8,9 @@ import (
 	"reflect"
 	"regexp"
 	"strings"
-	"time"
 
 	"github.com/pkg/errors"
+	goutils "go.viam.com/utils"
 
 	"go.viam.com/rdk/logging"
 	"go.viam.com/rdk/resource"
@@ -42,7 +42,11 @@ type Module struct {
 	Environment map[string]string `json:"env,omitempty"`
 
 	// FirstRunTimeout is the timeout duration for the first run script.
-	FirstRunTimeout time.Duration `json:"first_run_timeout,omitempty"`
+	// Setting this field to a zero value (e.g. 0 * [time.Second]) will
+	// set the first run timeout to the default value of 1 hour, which
+	// is equivalent to leaving this field unset. If you wish to set an
+	// immediate timeout you can set this field to a negative value.
+	FirstRunTimeout goutils.Duration `json:"first_run_timeout,omitempty"`
 
 	// Status refers to the validations done in the APP to make sure a module is configured correctly
 	Status           *AppValidationStatus `json:"status"`

diff --git a/config/proto_conversions.go b/config/proto_conversions.go
@@ -9,6 +9,7 @@ import (
 	"github.com/pkg/errors"
 	packagespb "go.viam.com/api/app/packages/v1"
 	pb "go.viam.com/api/app/v1"
+	"go.viam.com/utils"
 	"go.viam.com/utils/pexec"
 	"go.viam.com/utils/protoutils"
 	"go.viam.com/utils/rpc"
@@ -309,7 +310,7 @@ func ModuleConfigToProto(module *Module) (*pb.ModuleConfig, error) {
 		ModuleId:        module.ModuleID,
 		Env:             module.Environment,
 		Status:          status,
-		FirstRunTimeout: durationpb.New(module.FirstRunTimeout),
+		FirstRunTimeout: durationpb.New(module.FirstRunTimeout.Unwrap()),
 	}
 
 	return &proto, nil
@@ -330,7 +331,7 @@ func ModuleConfigFromProto(proto *pb.ModuleConfig) (*Module, error) {
 		ModuleID:        proto.GetModuleId(),
 		Environment:     proto.GetEnv(),
 		Status:          status,
-		FirstRunTimeout: proto.GetFirstRunTimeout().AsDuration(),
+		FirstRunTimeout: utils.Duration(proto.GetFirstRunTimeout().AsDuration()),
 	}
 	return &module, nil
 }

diff --git a/config/proto_conversions_test.go b/config/proto_conversions_test.go
@@ -14,6 +14,7 @@ import (
 	packagespb "go.viam.com/api/app/packages/v1"
 	pb "go.viam.com/api/app/v1"
 	"go.viam.com/test"
+	goutils "go.viam.com/utils"
 	"go.viam.com/utils/jwks"
 	"go.viam.com/utils/pexec"
 	"go.viam.com/utils/rpc"
@@ -207,7 +208,7 @@ var testModuleWithTimeout = Module{
 	Type:            ModuleTypeLocal,
 	ModuleID:        "a:b",
 	Environment:     map[string]string{"SOME_VAR": "value"},
-	FirstRunTimeout: time.Minute,
+	FirstRunTimeout: goutils.Duration(time.Minute),
 }
 
 var testPackageConfig = PackageConfig{

diff --git a/go.mod b/go.mod
@@ -84,7 +84,7 @@ require (
 	go.uber.org/zap v1.27.0
 	go.viam.com/api v0.1.351
 	go.viam.com/test v1.1.1-0.20220913152726-5da9916c08a2
-	go.viam.com/utils v0.1.110
+	go.viam.com/utils v0.1.112
 	goji.io v2.0.2+incompatible
 	golang.org/x/image v0.19.0
 	golang.org/x/mobile v0.0.0-20240112133503-c713f31d574b

diff --git a/go.sum b/go.sum
@@ -1641,8 +1641,8 @@ go.viam.com/api v0.1.351 h1:jNIpa7K1RzQzWrwnXJBwFZ+oQ/OiZBWLqTp9+7nznPI=
 go.viam.com/api v0.1.351/go.mod h1:5lpVRxMsKFCaahqsnJfPGwJ9baoQ6PIKQu3lxvy6Wtw=
 go.viam.com/test v1.1.1-0.20220913152726-5da9916c08a2 h1:oBiK580EnEIzgFLU4lHOXmGAE3MxnVbeR7s1wp/F3Ps=
 go.viam.com/test v1.1.1-0.20220913152726-5da9916c08a2/go.mod h1:XM0tej6riszsiNLT16uoyq1YjuYPWlRBweTPRDanIts=
-go.viam.com/utils v0.1.110 h1:xV6rcJiNq4iKy1f7y/5JeoasEZM29314X7Mg8xSUKLk=
-go.viam.com/utils v0.1.110/go.mod h1:SYvcY/TKy9yv1d95era4IEehImkXffWu/5diDBS/4X4=
+go.viam.com/utils v0.1.112 h1:yuVkNITUijdP/CMI3BaDozUMZwP4Ari57BvRQfORFK0=
+go.viam.com/utils v0.1.112/go.mod h1:SYvcY/TKy9yv1d95era4IEehImkXffWu/5diDBS/4X4=
 go4.org/unsafe/assume-no-moving-gc v0.0.0-20230525183740-e7c30c78aeb2 h1:WJhcL4p+YeDxmZWg141nRm7XC8IDmhz7lk5GpadO1Sg=
 go4.org/unsafe/assume-no-moving-gc v0.0.0-20230525183740-e7c30c78aeb2/go.mod h1:FftLjUGFEDu5k8lt0ddY+HcrH/qU/0qk+H8j9/nTl3E=
 gocv.io/x/gocv v0.25.0/go.mod h1:Rar2PS6DV+T4FL+PM535EImD/h13hGVaHhnCu1xarBs=

diff --git a/module/modmanager/manager.go b/module/modmanager/manager.go
@@ -1111,8 +1111,11 @@ func (mgr *Manager) FirstRun(ctx context.Context, conf config.Module) error {
 
 	moduleEnvironment := getFullEnvironment(conf, dataDir, mgr.viamHomeDir)
 
-	// TODO(RSDK-9060): support a user-supplied timeout
-	cmdCtx, cancel := context.WithTimeout(ctx, defaultFirstRunTimeout)
+	timeout := defaultFirstRunTimeout
+	if conf.FirstRunTimeout > 0 {
+		timeout = conf.FirstRunTimeout.Unwrap()
+	}
+	cmdCtx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
 	//nolint:gosec // Yes, we are deliberating executing arbitrary user code here.