From 74cb28ec30acd07a86f1392b71c1ecc2afbf88ae Mon Sep 17 00:00:00 2001
From: Kris Hicks <khicks@hashicorp.com>
Date: Thu, 10 Dec 2020 07:27:15 -0800
Subject: [PATCH] pluginmanager: WaitForFirstFingerprint times out (#9597)

As pointed out by @tgross[1], prior to this change we would have been blocking
until all managers waited for first fingerprint rather than timing out as
intended.

1: https://github.com/hashicorp/nomad/pull/9590#discussion_r539534906
---
 client/pluginmanager/group.go      |  6 ++--
 client/pluginmanager/group_test.go | 57 ++++++++++++++++++++++++++++++
 client/pluginmanager/testing.go    | 16 +++++++--
 3 files changed, 74 insertions(+), 5 deletions(-)

diff --git a/client/pluginmanager/group.go b/client/pluginmanager/group.go
index 3448248ff..f6919738e 100644
--- a/client/pluginmanager/group.go
+++ b/client/pluginmanager/group.go
@@ -65,12 +65,12 @@ func (m *PluginGroup) WaitForFirstFingerprint(ctx context.Context) (<-chan struc
 		go func() {
 			defer wg.Done()
 			logger.Debug("waiting on plugin manager initial fingerprint")
-			<-manager.WaitForFirstFingerprint(ctx)
+
 			select {
+			case <-manager.WaitForFirstFingerprint(ctx):
+				logger.Debug("finished plugin manager initial fingerprint")
 			case <-ctx.Done():
 				logger.Warn("timeout waiting for plugin manager to be ready")
-			default:
-				logger.Debug("finished plugin manager initial fingerprint")
 			}
 		}()
 	}
diff --git a/client/pluginmanager/group_test.go b/client/pluginmanager/group_test.go
index 853308a15..07448d542 100644
--- a/client/pluginmanager/group_test.go
+++ b/client/pluginmanager/group_test.go
@@ -1,8 +1,10 @@
 package pluginmanager
 
 import (
+	"context"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/nomad/helper/testlog"
 	"github.com/stretchr/testify/require"
@@ -62,3 +64,58 @@ func TestPluginGroup_Shutdown(t *testing.T) {
 
 	require.Error(group.RegisterAndRun(&MockPluginManager{}))
 }
+
+func TestPluginGroup_WaitForFirstFingerprint(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+
+	managerCh := make(chan struct{})
+	manager := &MockPluginManager{
+		RunF:                      func() {},
+		WaitForFirstFingerprintCh: managerCh,
+	}
+
+	// close immediately to beat the context timeout
+	close(managerCh)
+
+	group := New(testlog.HCLogger(t))
+	require.NoError(group.RegisterAndRun(manager))
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	groupCh, err := group.WaitForFirstFingerprint(ctx)
+	require.NoError(err)
+
+	select {
+	case <-groupCh:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected groupCh to be closed")
+	}
+}
+
+func TestPluginGroup_WaitForFirstFingerprint_Timeout(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+
+	managerCh := make(chan struct{})
+	manager := &MockPluginManager{
+		RunF:                      func() {},
+		WaitForFirstFingerprintCh: managerCh,
+	}
+
+	group := New(testlog.HCLogger(t))
+	require.NoError(group.RegisterAndRun(manager))
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
+	defer cancel()
+
+	groupCh, err := group.WaitForFirstFingerprint(ctx)
+
+	select {
+	case <-groupCh:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected groupCh to be closed due to context timeout")
+	}
+	require.NoError(err)
+}
diff --git a/client/pluginmanager/testing.go b/client/pluginmanager/testing.go
index acb378414..c92aa33ae 100644
--- a/client/pluginmanager/testing.go
+++ b/client/pluginmanager/testing.go
@@ -1,10 +1,22 @@
 package pluginmanager
 
+import "context"
+
 type MockPluginManager struct {
-	RunF      func()
-	ShutdownF func()
+	RunF                      func()
+	ShutdownF                 func()
+	WaitForFirstFingerprintCh <-chan struct{}
 }
 
 func (m *MockPluginManager) Run()               { m.RunF() }
 func (m *MockPluginManager) Shutdown()          { m.ShutdownF() }
 func (m *MockPluginManager) PluginType() string { return "mock" }
+func (m *MockPluginManager) WaitForFirstFingerprint(ctx context.Context) <-chan struct{} {
+	if m.WaitForFirstFingerprintCh != nil {
+		return m.WaitForFirstFingerprintCh
+	}
+
+	ch := make(chan struct{})
+	close(ch)
+	return ch
+}