Simplify node lock and various small changes

Co-authored-by: Fabian Kammel <fabian@kammel.dev>
Co-authored-by: Daniel Weiße <66256922+daniel-weisse@users.noreply.github.com>

bootstrapper/internal/clean/clean.go

@@ -0,0 +1,64 @@
+package clean
+
+import (
+	"sync"
+)
+
+type cleaner struct {
+	stoppers  []stopper
+	stopC     chan struct{}
+	startOnce sync.Once
+	wg        sync.WaitGroup
+}
+
+// New creates a new cleaner.
+func New(stoppers ...stopper) *cleaner {
+	res := &cleaner{
+		stoppers: stoppers,
+		stopC:    make(chan struct{}, 1),
+	}
+	res.wg.Add(1) // for the Start goroutine
+	return res
+}
+
+// With adds a new stopper to the cleaner.
+func (c *cleaner) With(stopper stopper) *cleaner {
+	c.stoppers = append(c.stoppers, stopper)
+	return c
+}
+
+// Start blocks until it receives a stop message, stops all services gracefully and returns.
+func (c *cleaner) Start() {
+	c.startOnce.Do(func() {
+		defer c.wg.Done()
+		// wait for the stop message
+		<-c.stopC
+
+		c.wg.Add(len(c.stoppers))
+		for _, stopItem := range c.stoppers {
+			go func(stopItem stopper) {
+				defer c.wg.Done()
+				stopItem.Stop()
+			}(stopItem)
+		}
+	})
+}
+
+// Clean initiates the cleanup but does not wait for it to complete.
+func (c *cleaner) Clean() {
+	// try to enqueue the stop message once
+	// if the channel is full, the message is dropped
+	select {
+	case c.stopC <- struct{}{}:
+	default:
+	}
+}
+
+// Done waits for the cleanup to complete.
+func (c *cleaner) Done() {
+	c.wg.Wait()
+}
+
+type stopper interface {
+	Stop()
+}

bootstrapper/internal/clean/clean_test.go

@@ -0,0 +1,98 @@
+package clean
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/goleak"
+)
+
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
+func TestNew(t *testing.T) {
+	assert := assert.New(t)
+
+	cleaner := New(&spyStopper{})
+	assert.NotNil(cleaner)
+	assert.NotEmpty(cleaner.stoppers)
+}
+
+func TestWith(t *testing.T) {
+	assert := assert.New(t)
+
+	cleaner := New().With(&spyStopper{})
+	assert.NotEmpty(cleaner.stoppers)
+}
+
+func TestClean(t *testing.T) {
+	assert := assert.New(t)
+
+	stopper := &spyStopper{}
+	cleaner := New(stopper)
+	go cleaner.Start()
+	cleaner.Clean()
+	cleaner.Done()
+	assert.Equal(int64(1), atomic.LoadInt64(&stopper.stopped))
+	// call again to make sure it doesn't panic or block or clean up again
+	cleaner.Clean()
+	assert.Equal(int64(1), atomic.LoadInt64(&stopper.stopped))
+}
+
+func TestCleanBeforeStart(t *testing.T) {
+	assert := assert.New(t)
+	// calling Clean before Start should work
+	stopper := &spyStopper{}
+	cleaner := New(stopper)
+	cleaner.Clean()
+	cleaner.Start()
+	cleaner.Done()
+	assert.Equal(int64(1), atomic.LoadInt64(&stopper.stopped))
+}
+
+func TestConcurrent(t *testing.T) {
+	assert := assert.New(t)
+	// calling Clean concurrently should call Stop exactly once
+
+	stopper := &spyStopper{}
+	cleaner := New(stopper)
+
+	parallelism := 10
+	wg := sync.WaitGroup{}
+
+	start := func() {
+		defer wg.Done()
+		cleaner.Start()
+	}
+
+	clean := func() {
+		defer wg.Done()
+		cleaner.Clean()
+	}
+
+	done := func() {
+		defer wg.Done()
+		cleaner.Done()
+	}
+
+	wg.Add(3 * parallelism)
+	for i := 0; i < parallelism; i++ {
+		go start()
+		go clean()
+		go done()
+	}
+	wg.Wait()
+	cleaner.Done()
+	assert.Equal(int64(1), atomic.LoadInt64(&stopper.stopped))
+}
+
+type spyStopper struct {
+	stopped int64
+}
+
+func (s *spyStopper) Stop() {
+	atomic.AddInt64(&s.stopped, 1)
+}