Sevey/shutdown management (#1434)

<!--
Please read and fill out this form before submitting your PR.

Please make sure you have reviewed our contributors guide before
submitting your
first PR.
-->

## Overview

<!-- 
Please provide an explanation of the PR, including the appropriate
context,
background, goal, and rationale. If there is an issue with this
information,
please provide a tl;dr and link the issue. 
-->

Context #1412 
Closes #1412 

Tested by running the test suite in a loop until the runner timeout with
no panics for writing to closed files:
https://github.com/rollkit/rollkit/actions/runs/7426536284/job/20210452949?pr=1434

## Checklist

<!-- 
Please complete the checklist to ensure that the PR is ready to be
reviewed.

IMPORTANT:
PRs should be left in Draft until the below checklist is completed.
-->

- [x] New and updated code has appropriate documentation
- [ ] New and updated code has new and/or updated testing
- [ ] Required CI checks are passing
- [ ] Visual proof for any user facing features like CLI or
documentation updates
- [x] Linked issues closed with keywords


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Implemented a new thread management system to enhance app performance
and stability.

- **Refactor**
- Improved internal thread handling mechanisms for better resource
management.

- **Tests**
- Added tests for the new thread management features to ensure
reliability.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: MSevey

node/full.go

@@ -34,6 +34,7 @@ import (
 	"github.com/rollkit/rollkit/state/txindex"
 	"github.com/rollkit/rollkit/state/txindex/kv"
 	"github.com/rollkit/rollkit/store"
+	"github.com/rollkit/rollkit/types"
 )
 
 // prefixes used in KV store to separate main node data from DALC data
@@ -82,8 +83,9 @@ type FullNode struct {
 
 	// keep context here only because of API compatibility
 	// - it's used in `OnStart` (defined in service.Service interface)
-	ctx    context.Context
-	cancel context.CancelFunc
+	ctx           context.Context
+	cancel        context.CancelFunc
+	threadManager *types.ThreadManager
 }
 
 // newFullNode creates a new Rollkit full node.
@@ -175,6 +177,7 @@ func newFullNode(
 		bSyncService:   blockSyncService,
 		ctx:            ctx,
 		cancel:         cancel,
+		threadManager:  types.NewThreadManager(),
 	}
 
 	node.BaseService = *service.NewBaseService(logger, "Node", node)
@@ -339,15 +342,15 @@ func (n *FullNode) OnStart() error {
 
 	if n.nodeConfig.Aggregator {
 		n.Logger.Info("working in aggregator mode", "block time", n.nodeConfig.BlockTime)
-		go n.blockManager.AggregationLoop(n.ctx, n.nodeConfig.LazyAggregator)
-		go n.blockManager.BlockSubmissionLoop(n.ctx)
-		go n.headerPublishLoop(n.ctx)
-		go n.blockPublishLoop(n.ctx)
+		n.threadManager.Go(func() { n.blockManager.AggregationLoop(n.ctx, n.nodeConfig.LazyAggregator) })
+		n.threadManager.Go(func() { n.blockManager.BlockSubmissionLoop(n.ctx) })
+		n.threadManager.Go(func() { n.headerPublishLoop(n.ctx) })
+		n.threadManager.Go(func() { n.blockPublishLoop(n.ctx) })
 		return nil
 	}
-	go n.blockManager.RetrieveLoop(n.ctx)
-	go n.blockManager.BlockStoreRetrieveLoop(n.ctx)
-	go n.blockManager.SyncLoop(n.ctx, n.cancel)
+	n.threadManager.Go(func() { n.blockManager.RetrieveLoop(n.ctx) })
+	n.threadManager.Go(func() { n.blockManager.BlockStoreRetrieveLoop(n.ctx) })
+	n.threadManager.Go(func() { n.blockManager.SyncLoop(n.ctx, n.cancel) })
 	return nil
 }
 
@@ -369,6 +372,8 @@ func (n *FullNode) GetGenesisChunks() ([]string, error) {
 func (n *FullNode) OnStop() {
 	n.Logger.Info("halting full node...")
 	n.cancel()
+	n.threadManager.Wait()
+	n.Logger.Info("shutting down full node sub services...")
 	err := n.p2pClient.Close()
 	err = multierr.Append(err, n.hSyncService.Stop())
 	err = multierr.Append(err, n.bSyncService.Stop())

types/threadmanager.go

@@ -0,0 +1,28 @@
+package types
+
+import "sync"
+
+// ThreadManager is a simple wrapper around sync.WaitGroup to make it easier to
+// manage threads
+type ThreadManager struct {
+	wg sync.WaitGroup
+}
+
+// NewThreadManager creates a new ThreadManager
+func NewThreadManager() *ThreadManager {
+	return &ThreadManager{}
+}
+
+// Wait blocks until all goroutines have called Done on the WaitGroup
+func (tm *ThreadManager) Wait() {
+	tm.wg.Wait()
+}
+
+// Go launches a goroutine and adds it to the WaitGroup
+func (tm *ThreadManager) Go(f func()) {
+	tm.wg.Add(1)
+	go func() {
+		defer tm.wg.Done()
+		f()
+	}()
+}

types/threadmanager_test.go

@@ -0,0 +1,30 @@
+package types
+
+import (
+	"testing"
+	"time"
+)
+
+func TestThreadManager(t *testing.T) {
+	tm := NewThreadManager()
+	start := time.Now()
+
+	// Launch a goroutine that waits for 10ms
+	tm.Go(func() {
+		time.Sleep(10 * time.Millisecond)
+	})
+	// Launch a goroutine that waits for 100ms
+	tm.Go(func() {
+		time.Sleep(100 * time.Millisecond)
+	})
+
+	// This will block until the goroutines have called Done on the
+	// WaitGroup
+	//
+	// This also implicitly verifies that we are properly calling Done(),
+	// otherwise this test would time out the test suite
+	tm.Wait()
+	if time.Since(start) < 100*time.Millisecond {
+		t.Errorf("WaitGroup returned before all goroutines are done")
+	}
+}