From 631796a7260e1e9a07dc06920b9877215a38aed2 Mon Sep 17 00:00:00 2001
From: Sasha Koshka <supergriffin64@gmail.com>
Date: Wed, 31 May 2023 18:07:38 -0400
Subject: [PATCH] Routine manager can now be shut down

---
 routines/routines.go | 66 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 62 insertions(+), 4 deletions(-)

diff --git a/routines/routines.go b/routines/routines.go
index f39c3f2..f727bb2 100644
--- a/routines/routines.go
+++ b/routines/routines.go
@@ -7,9 +7,43 @@ import "log"
 import "time"
 import "sync"
 
-// Routine is a long-running function that does not return until it is finished.
-// An error is returned if the routine exited due to an error.
-type Routine func () error
+type routine struct {
+	run, shutdown func () error
+}
+
+func (routine routine) Run () error {
+	if routine.run == nil {
+		return nil
+	} else {
+		return routine.run()
+	}
+}
+
+func (routine routine) Shutdown () error {
+	if routine.shutdown == nil {
+		return nil
+	} else {
+		return routine.shutdown()
+	}
+}
+
+// From creates a routine from a separate run and shutdown function.
+func From (run, shutdown func () error) Routine {
+	return routine {
+		run:      run,
+		shutdown: shutdown,
+	}
+} 
+
+// Routine is an object that can be run and stopped.
+type Routine interface {
+	// Run is a long-running function that does not return until it is
+	// finished. An error is returned if the routine exited due to an error.
+	Run () error
+
+	// Shutdown stops Run.
+	Shutdown () error
+}
 
 // Manager is a system capable of managing multiple routines, and restarting
 // them if they fail.
@@ -27,6 +61,9 @@ type Manager struct {
 	// is nil, messages will be written to the standard logger. To disable
 	// logging altogether, this can be set to io.Discard.
 	Logger io.Writer
+
+	stoppingMutex sync.Mutex
+	stopping bool
 }
 
 // Run spawns all routines in the Routines slice. If a routine exits with an
@@ -47,6 +84,21 @@ func (manager *Manager) Run () error {
 	return errExit
 }
 
+// Shutdown shuts down all routines in the manager.
+func (manager *Manager) Shutdown () (err error) {
+	manager.stoppingMutex.Lock()
+	manager.stopping = true
+	manager.stoppingMutex.Unlock()
+
+	for _, routine := range manager.Routines {
+		routineErr := routine.Shutdown()
+		if routineErr != nil {
+			err = routineErr
+		}
+	}
+	return
+}
+
 // Append adds one or more routines to the Routines slice. This has no effect if
 // the manager is already running.
 func (manager *Manager) Append (routines ...Routine) {
@@ -66,9 +118,15 @@ func (manager *Manager) runRoutine (routine Routine, group *sync.WaitGroup, errE
 
 	var err error
 	for {
+		stopping := false
+		manager.stoppingMutex.Lock()
+		stopping = manager.stopping
+		manager.stoppingMutex.Unlock()
+		if stopping { break }
+	
 		// TODO: recover from panics
 		lastStart := time.Now()
-		err = routine()
+		err = routine.Run()
 
 		if err == nil {
 			manager.log("(i) routine exited")