From 2a5adf39c9615c710af8196f068d984a40c65f5b Mon Sep 17 00:00:00 2001
From: Alexander Kratzsch <akratzsch@baffinbaynetworks.com>
Date: Fri, 28 Jan 2022 14:13:38 +0100
Subject: [PATCH] Implement InstallRequest State as StateMachine

---
 database.go      | 50 ++++++++++++++++++++++++-----
 state_machine.go | 82 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 124 insertions(+), 8 deletions(-)
 create mode 100644 state_machine.go

diff --git a/database.go b/database.go
index 0028e1e..aa8489e 100644
--- a/database.go
+++ b/database.go
@@ -56,10 +56,7 @@ func (m *Machine) State() string {
 		return "provisioned"
 	}
 
-	if m.installRequest.State == "" {
-		return "unknown"
-	}
-	return m.installRequest.State
+	return m.installRequest.State.Current()
 }
 
 // When calling this function, you should hold a read-lock on the db object
@@ -74,10 +71,45 @@ func (m *Machine) getInstallRequest(db *tateruDB) {
 	}
 }
 
+func NewInstallRequestStateMachine() *StateMachine {
+	states := StateMap{
+		"provisioned": {
+			Events: EventMap{
+				"RECEIVED_BOOT_INSTALLER_REQUEST": Event{
+					NewState: "pending",
+				},
+			},
+		},
+		"pending": {
+			Events: EventMap{
+				"SENT_BOOT_INSTALLER_REQUEST_TO_MANAGER": Event{
+					NewState: "booting",
+				},
+			},
+		},
+		"booting": {
+			Events: EventMap{
+				"RECEIVED_INSTALLER_CALLBACK": Event{
+					NewState: "booted",
+				},
+			},
+		},
+		"booted": {
+			Events: EventMap{
+				"RECEIVED_INSTALLER_EXITING_CALLBACK": Event{
+					NewState: "provisioned",
+				},
+			},
+		},
+	}
+
+	return NewStateMachine("pending", states)
+}
+
 type InstallRequest struct {
 	LastUpdate    time.Time
 	Nonce         string
-	State         string
+	State         *StateMachine
 	SSHPubKey     string
 	InstallerAddr string
 	SSHPorts      SSHPorts
@@ -292,7 +324,7 @@ func (db *tateruDB) HandleBootInstallerAPI(w http.ResponseWriter, r *http.Reques
 
 	installRequest := InstallRequest{
 		LastUpdate: time.Now(),
-		State:      "pending",
+		State:      NewInstallRequestStateMachine(),
 		SSHPubKey:  bir.SSHPubKey,
 		Nonce:      bir.Nonce,
 	}
@@ -330,9 +362,11 @@ func (db *tateruDB) HandleBootInstallerAPI(w http.ResponseWriter, r *http.Reques
 	}
 
 	installRequest.LastUpdate = time.Now()
-	installRequest.State = "booting"
+	installRequest.State.Transition("SENT_BOOT_INSTALLER_REQUEST_TO_MANAGER")
 	db.installRequests[uuid] = installRequest
 
+	installRequest.State.WaitFor("booted")
+
 	return
 }
 
@@ -397,7 +431,7 @@ func (db *tateruDB) HandleInstallerCallbackAPI(w http.ResponseWriter, r *http.Re
 	}
 
 	installRequest.LastUpdate = time.Now()
-	installRequest.State = "booted"
+	installRequest.State.Transition("RECEIVED_INSTALLER_CALLBACK")
 	db.installRequests[uuid] = installRequest
 
 	w.Header().Add("content-type", "application/json; charset=utf-8")
diff --git a/state_machine.go b/state_machine.go
new file mode 100644
index 0000000..acaaa8a
--- /dev/null
+++ b/state_machine.go
@@ -0,0 +1,82 @@
+package main
+
+import (
+	"errors"
+	"sync"
+)
+
+type EventName string
+
+type Event struct {
+	NewState StateName
+}
+
+type EventMap map[EventName]Event
+
+type StateName string
+
+type State struct {
+	Events EventMap
+}
+
+type StateMap map[StateName]State
+
+type StateMachine struct {
+	InitialState StateName
+	States       StateMap
+
+	current StateName
+	cond    *sync.Cond
+}
+
+func NewStateMachine(initialState StateName, states StateMap) *StateMachine {
+	return &StateMachine{
+		InitialState: initialState,
+		States:       states,
+		cond:         sync.NewCond(&sync.Mutex{}),
+	}
+}
+
+func (s *StateMachine) Current() StateName {
+	s.cond.L.Lock()
+	current := s.current
+	s.cond.L.Unlock()
+
+	if current == "" {
+		current = s.InitialState
+	}
+
+	return current
+}
+
+func (s *StateMachine) String() string {
+	return string(s.Current())
+}
+
+func (s *StateMachine) Transition(event EventName) (StateName, error) {
+	current := s.Current()
+
+	transitions := s.States[current].Events
+	transition, ok := transitions[event]
+	if !ok {
+		return "", errors.New("Invalid transition: Event not available for current state")
+	}
+
+	if transition.NewState != "" {
+		s.cond.L.Lock()
+		s.current = transition.NewState
+		s.cond.Broadcast()
+		s.cond.L.Unlock()
+	}
+
+	return s.Current(), nil
+}
+
+func (s *StateMachine) WaitFor(state StateName) {
+	s.cond.L.Lock()
+	for s.current != state {
+		s.cond.Wait()
+	}
+	s.cond.L.Unlock()
+	return
+}