pkg/agent: improve reacting to shutdown signal when waiting for ok to reboot

Previous attempt in 890518e did not
make it responsive enough and actually introduced a bug as we always log
the error now as a warning, even if it's nil.

Making this more efficient actually unveiled one more bug in tests code,
which has also been fixed.

Signed-off-by: Mateusz Gozdek <[email protected]>

Author: invidian

pkg/agent/agent.go

@@ -209,23 +209,30 @@ func (k *klocksmith) process(ctx context.Context) error {
 	go k.watchUpdateStatus(ctx, k.updateStatusCallback)
 
 	// Block until constants.AnnotationOkToReboot is set.
-	okToReboot := false
-	for !okToReboot {
+	for okToReboot := false; !okToReboot; {
+		klog.Infof("Waiting for ok-to-reboot from controller...")
+
+		errCh := make(chan error)
+
+		go func() {
+			errCh <- k.waitForOkToReboot(ctx)
+		}()
+
 		select {
 		case <-ctx.Done():
 			klog.Infof("Got stop signal while waiting for ok-to-reboot from controller")
 
 			return nil
-		default:
-			klog.Infof("Waiting for ok-to-reboot from controller...")
+		case err := <-errCh:
+			if err != nil {
+				klog.Warningf("Error waiting for an ok-to-reboot: %v", err)
 
-			err := k.waitForOkToReboot(ctx)
-			if err == nil {
-				// Time to reboot.
-				okToReboot = true
+				// Break select statement to restart watching for ok to reboot.
+				break
 			}
 
-			klog.Warningf("Error waiting for an ok-to-reboot: %v", err)
+			// Time to reboot.
+			okToReboot = true
 		}
 	}
 

pkg/agent/agent_test.go

@@ -1549,8 +1549,6 @@ func Test_Running_agent(t *testing.T) {
 		t.Run("agent_receives_termination_signal_while_waiting_for_all_pods_to_be_terminated", func(t *testing.T) {
 			t.Parallel()
 
-			rebootTriggerred := make(chan bool, 1)
-
 			podsToCreate := []*corev1.Pod{
 				{
 					ObjectMeta: metav1.ObjectMeta{
@@ -1568,6 +1566,19 @@ func Test_Running_agent(t *testing.T) {
 			fakeClient := fake.NewSimpleClientset(podsToCreate[0], testNode())
 			addEvictionSupport(t, fakeClient)
 
+			podDeletionAttemptCh := make(chan struct{}, 1)
+
+			actionF := func(action k8stesting.Action) (bool, runtime.Object, error) {
+				podDeletionAttemptCh <- struct{}{}
+
+				// Never actually remove any pod.
+				return true, nil, nil
+			}
+
+			fakeClient.PrependReactor("create", "pods/eviction", actionF)
+
+			rebootTriggerred := make(chan bool, 1)
+
 			testConfig, node, _ := validTestConfig(t, testNode())
 			testConfig.Clientset = fakeClient
 			testConfig.PodDeletionGracePeriod = 30 * time.Second
@@ -1589,6 +1600,9 @@ func Test_Running_agent(t *testing.T) {
 
 			okToReboot(ctx, t, testConfig.Clientset.CoreV1().Nodes(), node.Name)
 
+			// Wait until we try to delete a pod.
+			<-podDeletionAttemptCh
+
 			cancel()
 
 			select {
@@ -2017,6 +2031,10 @@ func addEvictionSupport(t *testing.T, clientset *fake.Clientset) {
 		GroupVersion: "policy/v1",
 	}
 	clientset.Resources = append(clientset.Resources, coreResources, policyResources)
+
+	clientset.PrependReactor("create", "pods", func(action k8stesting.Action) (bool, runtime.Object, error) {
+		return action.GetSubresource() == "eviction", nil, nil
+	})
 }
 
 func testPodControllerReference() []metav1.OwnerReference {