From 9d658e97cfd60e113b55fcf663d5e1236cb4cb93 Mon Sep 17 00:00:00 2001
From: Rebecca Mahany-Horton <rebeccamahany@gmail.com>
Date: Wed, 11 Oct 2023 17:20:01 -0400
Subject: [PATCH] Allow power event watcher interrupt to be called multiple
 times

---
 .../power_event_watcher_other.go              | 10 +++-
 .../power_event_watcher_other_test.go         | 53 +++++++++++++++++++
 .../power_event_watcher_windows.go            |  8 +++
 3 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 pkg/windows/powereventwatcher/power_event_watcher_other_test.go

diff --git a/pkg/windows/powereventwatcher/power_event_watcher_other.go b/pkg/windows/powereventwatcher/power_event_watcher_other.go
index 1325aa460..94f41142e 100644
--- a/pkg/windows/powereventwatcher/power_event_watcher_other.go
+++ b/pkg/windows/powereventwatcher/power_event_watcher_other.go
@@ -9,7 +9,8 @@ import (
 )
 
 type noOpPowerEventWatcher struct {
-	interrupt chan struct{}
+	interrupt   chan struct{}
+	interrupted bool
 }
 
 func New(_ types.Knapsack, _ log.Logger) (*noOpPowerEventWatcher, error) {
@@ -24,5 +25,12 @@ func (n *noOpPowerEventWatcher) Execute() error {
 }
 
 func (n *noOpPowerEventWatcher) Interrupt(_ error) {
+	// Only perform shutdown tasks on first call to interrupt -- no need to repeat on potential extra calls.
+	if n.interrupted {
+		return
+	}
+
+	n.interrupted = true
+
 	n.interrupt <- struct{}{}
 }
diff --git a/pkg/windows/powereventwatcher/power_event_watcher_other_test.go b/pkg/windows/powereventwatcher/power_event_watcher_other_test.go
new file mode 100644
index 000000000..8349b8fb2
--- /dev/null
+++ b/pkg/windows/powereventwatcher/power_event_watcher_other_test.go
@@ -0,0 +1,53 @@
+//go:build !windows
+// +build !windows
+
+package powereventwatcher
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/go-kit/kit/log"
+	typesmocks "github.com/kolide/launcher/pkg/agent/types/mocks"
+	"github.com/stretchr/testify/require"
+)
+
+func TestInterrupt_Multiple(t *testing.T) {
+	t.Parallel()
+
+	p, err := New(typesmocks.NewKnapsack(t), log.NewNopLogger())
+	require.NoError(t, err)
+
+	// Start and then interrupt
+	go p.Execute()
+	p.Interrupt(errors.New("test error"))
+
+	// Confirm we can call Interrupt multiple times without blocking
+	interruptComplete := make(chan struct{})
+	expectedInterrupts := 3
+	for i := 0; i < expectedInterrupts; i += 1 {
+		go func() {
+			p.Interrupt(nil)
+			interruptComplete <- struct{}{}
+		}()
+	}
+
+	receivedInterrupts := 0
+	for {
+		if receivedInterrupts >= expectedInterrupts {
+			break
+		}
+
+		select {
+		case <-interruptComplete:
+			receivedInterrupts += 1
+			continue
+		case <-time.After(5 * time.Second):
+			t.Errorf("could not call interrupt multiple times and return within 5 seconds -- received %d interrupts before timeout", receivedInterrupts)
+			t.FailNow()
+		}
+	}
+
+	require.Equal(t, expectedInterrupts, receivedInterrupts)
+}
diff --git a/pkg/windows/powereventwatcher/power_event_watcher_windows.go b/pkg/windows/powereventwatcher/power_event_watcher_windows.go
index 9c1ca183c..44b81cded 100644
--- a/pkg/windows/powereventwatcher/power_event_watcher_windows.go
+++ b/pkg/windows/powereventwatcher/power_event_watcher_windows.go
@@ -33,6 +33,7 @@ type (
 		unsubscribeProcedure    *syscall.LazyProc
 		renderEventLogProcedure *syscall.LazyProc
 		interrupt               chan struct{}
+		interrupted             bool
 	}
 )
 
@@ -102,6 +103,13 @@ func (p *powerEventWatcher) Execute() error {
 }
 
 func (p *powerEventWatcher) Interrupt(_ error) {
+	// Only perform shutdown tasks on first call to interrupt -- no need to repeat on potential extra calls.
+	if p.interrupted {
+		return
+	}
+
+	p.interrupted = true
+
 	// EvtClose: https://learn.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtclose
 	ret, _, err := p.unsubscribeProcedure.Call(p.subscriptionHandle)
 	level.Debug(p.logger).Log("msg", "unsubscribed from power events", "ret", fmt.Sprintf("%+v", ret), "last_err", err)