From 816c3e1fc3d87c1a0e4594b5fe638e1caac14630 Mon Sep 17 00:00:00 2001
From: Sasha Koshka <sashakoshka@tebibyte.media>
Date: Sun, 3 Nov 2024 17:20:58 -0500
Subject: [PATCH] routines: Add system for routine initialization

---
 routines/routines.go | 42 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/routines/routines.go b/routines/routines.go
index 233ceb7..0b67e1c 100644
--- a/routines/routines.go
+++ b/routines/routines.go
@@ -26,6 +26,16 @@ type Routine interface {
 	Run (context.Context) error
 }
 
+// InitRoutine is a routine that has to be initialized before it can be run.
+type InitRoutine interface {
+	Routine
+
+	// Init is an initialization function that is called before any routines
+	// are run, including this one. Routines must not depend on eachother in
+	// this function.
+	Init (context.Context) error
+}
+
 // Manager is a system capable of managing multiple routines, and restarting
 // them if they fail.
 type Manager struct {
@@ -41,9 +51,13 @@ type Manager struct {
 	// Logger, if non-nil, is where log messages will be written to. If it
 	// is nil, messages will be written to the standard logger. To disable
 	// logging altogether, this can be set to io.Discard.
-	Logger io.Writer
+	Logger     io.Writer
+	loggerLock sync.Mutex
 
 	ctx context.Context
+
+	failed     map[Routine] struct { }
+	failedLock sync.Mutex
 }
 
 // Run spawns all routines in the Routines slice. If a routine exits with an
@@ -52,14 +66,24 @@ type Manager struct {
 func (this *Manager) Run (ctx context.Context) error {
 	ctx, done := context.WithCancel(ctx)
 	this.ctx = ctx
+	this.failed = make(map[Routine] struct { })
 
 	var waitGroup sync.WaitGroup
 	for _, routine := range this.Routines {
 		if routine != nil {
+			waitGroup.Add(1)
+			go this.initRoutine(routine, &waitGroup)
+		}
+	}
+	waitGroup.Wait()
+	
+	for _, routine := range this.Routines {
+		if _, failed := this.failed[routine]; !failed {
 			waitGroup.Add(1)
 			go this.runRoutine(routine, &waitGroup)
 		}
 	}
+	this.failed = nil
 	waitGroup.Wait()
 	
 	done()
@@ -73,6 +97,8 @@ func (this *Manager) Append (routines ...Routine) {
 }
 
 func (this *Manager) log (message ...any) {
+	this.loggerLock.Lock()
+	defer this.loggerLock.Unlock()
 	if this.Logger == nil {
 		log.Println(message...)
 	} else {
@@ -80,6 +106,20 @@ func (this *Manager) log (message ...any) {
 	}
 }
 
+func (this *Manager) initRoutine (routine Routine, group *sync.WaitGroup) {
+	defer group.Done()
+
+	if initRoutine, ok := routine.(InitRoutine); ok {
+		err := initRoutine.Init(this.ctx)
+		if err != nil {
+			this.log("XXX routine failed to initialize:", err)
+			this.failedLock.Lock()
+			defer this.failedLock.Unlock()
+			this.failed[routine] = struct { } { }
+		}
+	}
+}
+
 func (this *Manager) runRoutine (routine Routine, group *sync.WaitGroup) {
 	defer group.Done()