From 5cbd01329f18bcc9b66d02c3e0bba5ff88f26fdd Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Wed, 12 Sep 2018 09:36:04 -0700
Subject: [PATCH] health_hook: fix panic and add tests

Still more testing to do, but I want to get this panic fixed ASAP.

All new tests pass with -race
---
 client/allocrunnerv2/health_hook.go      |   1 +
 client/allocrunnerv2/health_hook_test.go | 170 +++++++++++++++++++++++
 2 files changed, 171 insertions(+)
 create mode 100644 client/allocrunnerv2/health_hook_test.go

diff --git a/client/allocrunnerv2/health_hook.go b/client/allocrunnerv2/health_hook.go
index 43c6ef374..1b51fdbb3 100644
--- a/client/allocrunnerv2/health_hook.go
+++ b/client/allocrunnerv2/health_hook.go
@@ -103,6 +103,7 @@ func newAllocHealthWatcherHook(logger log.Logger, alloc *structs.Allocation, hs
 		consul:       consul,
 		healthSetter: hs,
 		listener:     listener,
+		watchLock:    newWatchSema(),
 	}
 
 	h.logger = logger.Named(h.Name())
diff --git a/client/allocrunnerv2/health_hook_test.go b/client/allocrunnerv2/health_hook_test.go
new file mode 100644
index 000000000..42b2d4ac5
--- /dev/null
+++ b/client/allocrunnerv2/health_hook_test.go
@@ -0,0 +1,170 @@
+package allocrunnerv2
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/nomad/client/allocrunnerv2/interfaces"
+	"github.com/hashicorp/nomad/client/consul"
+	cstructs "github.com/hashicorp/nomad/client/structs"
+	"github.com/hashicorp/nomad/helper/testlog"
+	"github.com/hashicorp/nomad/nomad/mock"
+	"github.com/hashicorp/nomad/nomad/structs"
+	"github.com/stretchr/testify/require"
+)
+
+// statically assert health hook implements the expected interfaces
+var _ interfaces.RunnerPrerunHook = (*allocHealthWatcherHook)(nil)
+var _ interfaces.RunnerUpdateHook = (*allocHealthWatcherHook)(nil)
+var _ interfaces.RunnerDestroyHook = (*allocHealthWatcherHook)(nil)
+
+// mockHealthSetter implements healthSetter that stores health internally
+type mockHealthSetter struct {
+	setCalls   int
+	clearCalls int
+	healthy    *bool
+	isDeploy   *bool
+	taskEvents map[string]*structs.TaskEvent
+	mu         sync.Mutex
+}
+
+func (m *mockHealthSetter) SetHealth(healthy, isDeploy bool, taskEvents map[string]*structs.TaskEvent) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.setCalls++
+	m.healthy = &healthy
+	m.isDeploy = &isDeploy
+	m.taskEvents = taskEvents
+}
+
+func (m *mockHealthSetter) ClearHealth() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.clearCalls++
+	m.healthy = nil
+	m.isDeploy = nil
+	m.taskEvents = nil
+}
+
+func TestHealthHook_PrerunDestroy(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+
+	b := cstructs.NewAllocBroadcaster()
+	defer b.Close()
+
+	consul := consul.NewMockConsulServiceClient(t)
+	hs := &mockHealthSetter{}
+
+	h := newAllocHealthWatcherHook(testlog.HCLogger(t), mock.Alloc(), hs, b.Listen(), consul)
+
+	// Assert we implemented the right interfaces
+	prerunh, ok := h.(interfaces.RunnerPrerunHook)
+	require.True(ok)
+	_, ok = h.(interfaces.RunnerUpdateHook)
+	require.True(ok)
+	destroyh, ok := h.(interfaces.RunnerDestroyHook)
+	require.True(ok)
+
+	// Prerun
+	require.NoError(prerunh.Prerun(context.Background()))
+
+	// Destroy
+	require.NoError(destroyh.Destroy())
+}
+
+// TestHealthHook_SystemNoop asserts that system jobs return the noop tracker.
+func TestHealthHook_SystemNoop(t *testing.T) {
+	t.Parallel()
+
+	h := newAllocHealthWatcherHook(testlog.HCLogger(t), mock.SystemAlloc(), nil, nil, nil)
+
+	// Assert that it's the noop impl
+	_, ok := h.(noopAllocHealthWatcherHook)
+	require.True(t, ok)
+
+	// Assert the noop impl does not implement any hooks
+	_, ok = h.(interfaces.RunnerPrerunHook)
+	require.False(t, ok)
+	_, ok = h.(interfaces.RunnerUpdateHook)
+	require.False(t, ok)
+	_, ok = h.(interfaces.RunnerDestroyHook)
+	require.False(t, ok)
+}
+
+// TestHealthHook_BatchNoop asserts that batch jobs return the noop tracker.
+func TestHealthHook_BatchNoop(t *testing.T) {
+	t.Parallel()
+
+	h := newAllocHealthWatcherHook(testlog.HCLogger(t), mock.BatchAlloc(), nil, nil, nil)
+
+	// Assert that it's the noop impl
+	_, ok := h.(noopAllocHealthWatcherHook)
+	require.True(t, ok)
+}
+
+// TestHealthHook_WatchSema asserts only one caller can acquire the watchSema
+// lock at once and all other callers are blocked.
+func TestHealthHook_WatchSema(t *testing.T) {
+	t.Parallel()
+
+	s := newWatchSema()
+
+	s.acquire()
+
+	// Second acquire should block
+	acq2 := make(chan struct{})
+	go func() {
+		s.acquire()
+		close(acq2)
+	}()
+
+	select {
+	case <-acq2:
+		t.Fatalf("should not have been able to acquire lock")
+	case <-time.After(33 * time.Millisecond):
+		// Ok! It's blocking!
+	}
+
+	// Release and assert second acquire is now active
+	s.release()
+
+	select {
+	case <-acq2:
+		// Ok! Second acquire unblocked!
+	case <-time.After(33 * time.Millisecond):
+		t.Fatalf("second acquire should not have blocked")
+	}
+
+	// wait() should block until lock is released
+	waitCh := make(chan struct{})
+	go func() {
+		s.wait()
+		close(waitCh)
+	}()
+
+	select {
+	case <-waitCh:
+		t.Fatalf("should not have been able to acquire lock")
+	case <-time.After(33 * time.Millisecond):
+		// Ok! It's blocking!
+	}
+
+	// Release and assert wait() finished
+	s.release()
+
+	select {
+	case <-waitCh:
+		// Ok! wait() unblocked!
+	case <-time.After(33 * time.Millisecond):
+		t.Fatalf("wait should not have blocked")
+	}
+
+	// Assert wait() did not acquire the lock
+	s.acquire()
+	s.release()
+}