From 837dd922e9ffb02fa4ff59f5e081329d2d60ad38 Mon Sep 17 00:00:00 2001
From: Cheuk <90270663+cheukt@users.noreply.github.com>
Date: Mon, 2 Oct 2023 12:34:20 -0400
Subject: [PATCH] RSDK-5220 - Early exit reconfiguration on context.Done()
 (#3020)

---
 robot/impl/local_robot.go            | 22 +++++++-
 robot/impl/resource_manager.go       | 11 +++-
 robot/impl/robot_reconfigure_test.go | 77 ++++++++++++++++++++++++++++
 3 files changed, 107 insertions(+), 3 deletions(-)

diff --git a/robot/impl/local_robot.go b/robot/impl/local_robot.go
index 7d4a4eeb743..a76087ffbb6 100644
--- a/robot/impl/local_robot.go
+++ b/robot/impl/local_robot.go
@@ -721,11 +721,20 @@ func (r *localRobot) updateWeakDependents(ctx context.Context) {
 		select {
 		case <-resChan:
 		case <-ctxWithTimeout.Done():
-			r.logger.Warn(resource.NewBuildTimeoutError(resName))
+			if errors.Is(ctxWithTimeout.Err(), context.DeadlineExceeded) {
+				r.logger.Warn(resource.NewBuildTimeoutError(resName))
+			}
+		case <-ctx.Done():
+			return
 		}
 	}
 
 	for resName, res := range internalResources {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
 		resChan := make(chan struct{}, 1)
 		resName := resName
 		res := res
@@ -758,6 +767,11 @@ func (r *localRobot) updateWeakDependents(ctx context.Context) {
 
 	cfg := r.Config()
 	for _, conf := range append(cfg.Components, cfg.Services...) {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
 		conf := conf
 		ctxWithTimeout, timeoutCancel := context.WithTimeout(ctx, timeout)
 		defer timeoutCancel()
@@ -773,7 +787,11 @@ func (r *localRobot) updateWeakDependents(ctx context.Context) {
 		select {
 		case <-resChan:
 		case <-ctxWithTimeout.Done():
-			r.logger.Warn(resource.NewBuildTimeoutError(conf.ResourceName()))
+			if errors.Is(ctxWithTimeout.Err(), context.DeadlineExceeded) {
+				r.logger.Warn(resource.NewBuildTimeoutError(conf.ResourceName()))
+			}
+		case <-ctx.Done():
+			return
 		}
 	}
 }
diff --git a/robot/impl/resource_manager.go b/robot/impl/resource_manager.go
index fe2fd80380c..8c11cd03d0f 100644
--- a/robot/impl/resource_manager.go
+++ b/robot/impl/resource_manager.go
@@ -531,6 +531,11 @@ func (manager *resourceManager) completeConfig(
 	resourceNames := manager.resources.ReverseTopologicalSort()
 	timeout := rutils.GetResourceConfigurationTimeout(manager.logger)
 	for _, resName := range resourceNames {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
 		resChan := make(chan struct{}, 1)
 		resName := resName
 		ctxWithTimeout, timeoutCancel := context.WithTimeout(ctx, timeout)
@@ -605,7 +610,11 @@ func (manager *resourceManager) completeConfig(
 		select {
 		case <-resChan:
 		case <-ctxWithTimeout.Done():
-			robot.logger.Warn(resource.NewBuildTimeoutError(resName))
+			if errors.Is(ctxWithTimeout.Err(), context.DeadlineExceeded) {
+				robot.logger.Warn(resource.NewBuildTimeoutError(resName))
+			}
+		case <-ctx.Done():
+			return
 		}
 	}
 }
diff --git a/robot/impl/robot_reconfigure_test.go b/robot/impl/robot_reconfigure_test.go
index dba4dee7b7b..ee42bf8c9ee 100644
--- a/robot/impl/robot_reconfigure_test.go
+++ b/robot/impl/robot_reconfigure_test.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -3698,6 +3699,82 @@ func TestResourceConstructTimeout(t *testing.T) {
 	rr.reconfigureWorkers.Wait()
 }
 
+func TestResourceConstructCtxCancel(t *testing.T) {
+	logger := golog.NewTestLogger(t)
+
+	contructCount := 0
+	var wg sync.WaitGroup
+
+	mockAPI := resource.APINamespaceRDK.WithComponentType("mock")
+	modelName1 := utils.RandomAlphaString(5)
+	model1 := resource.DefaultModelFamily.WithModel(modelName1)
+
+	type cancelFunc struct {
+		c context.CancelFunc
+	}
+	var cFunc cancelFunc
+
+	resource.RegisterComponent(mockAPI, model1, resource.Registration[resource.Resource, resource.NoNativeConfig]{
+		Constructor: func(
+			ctx context.Context,
+			deps resource.Dependencies,
+			conf resource.Config,
+			logger golog.Logger,
+		) (resource.Resource, error) {
+			contructCount++
+			wg.Add(1)
+			defer wg.Done()
+			cFunc.c()
+			<-ctx.Done()
+			return &mockFake{Named: conf.ResourceName().AsNamed()}, nil
+		},
+	})
+	defer func() {
+		resource.Deregister(mockAPI, model1)
+	}()
+
+	cfg := &config.Config{
+		Components: []resource.Config{
+			{
+				Name:  "one",
+				Model: model1,
+				API:   mockAPI,
+			},
+			{
+				Name:  "two",
+				Model: model1,
+				API:   mockAPI,
+			},
+		},
+	}
+	t.Run("new", func(t *testing.T) {
+		contructCount = 0
+		ctxWithCancel, cancel := context.WithCancel(context.Background())
+		cFunc.c = cancel
+		r, err := New(ctxWithCancel, cfg, logger)
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, r.Close(context.Background()), test.ShouldBeNil)
+
+		wg.Wait()
+		test.That(t, contructCount, test.ShouldEqual, 1)
+	})
+	t.Run("reconfigure", func(t *testing.T) {
+		contructCount = 0
+		r, err := New(context.Background(), &config.Config{}, logger)
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, contructCount, test.ShouldEqual, 0)
+
+		ctxWithCancel, cancel := context.WithCancel(context.Background())
+		cFunc.c = cancel
+		r.Reconfigure(ctxWithCancel, cfg)
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, r.Close(context.Background()), test.ShouldBeNil)
+
+		wg.Wait()
+		test.That(t, contructCount, test.ShouldEqual, 1)
+	})
+}
+
 type mockFake struct {
 	resource.Named
 	createdAt      int