From f165e9e744b98eeae3f7175aee784d3ef20dd8f5 Mon Sep 17 00:00:00 2001
From: Andre Henriques <andr3h3nriqu3s@gmail.com>
Date: Tue, 16 Apr 2024 18:19:43 +0100
Subject: [PATCH] on service restart clean up mid way tasks closes #95

---
 config.toml           |  1 +
 logic/utils/config.go | 57 ++++++++++++++++++++++++++++++++++++++++---
 main.go               |  9 +------
 3 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/config.toml b/config.toml
index 494654a..86ba94c 100644
--- a/config.toml
+++ b/config.toml
@@ -5,6 +5,7 @@ HOSTNAME = "https://testing.andr3h3nriqu3s.com"
 NUMBER_OF_WORKERS = 20
 
 SUPRESS_CUDA = 1
+CLEAN_UP_ON_STARTUP = 1
 
 [ServiceUser]
 USER = "service"
diff --git a/logic/utils/config.go b/logic/utils/config.go
index 4788309..4cbc7d5 100644
--- a/logic/utils/config.go
+++ b/logic/utils/config.go
@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"
+	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/tasks/utils"
 
 	toml "github.com/BurntSushi/toml"
 	"github.com/charmbracelet/log"
@@ -29,7 +30,9 @@ type Config struct {
 	Hostname        string
 	Port            int
 	NumberOfWorkers int `toml:"number_of_workers"`
-	SupressCuda     int `toml:"supress_cuda"`
+
+	SupressCuda      int `toml:"supress_cuda"`
+	CleanUpOnStartup int `toml:"clean_up_on_startup"`
 
 	GpuWorker WorkerConfig `toml:"Worker"`
 
@@ -47,9 +50,11 @@ func LoadConfig() Config {
 		log.Error("Failed to load config file", "err", err)
 		// Use default values
 		return Config{
-			Hostname:        "localhost",
-			Port:            8000,
-			NumberOfWorkers: 10,
+			Hostname:         "localhost",
+			Port:             8000,
+			NumberOfWorkers:  10,
+			CleanUpOnStartup: 1,
+			SupressCuda:      1,
 			GpuWorker: WorkerConfig{
 				NumberOfWorkers: 1,
 				Pulling:         "500ms",
@@ -76,6 +81,50 @@ func LoadConfig() Config {
 	return conf
 }
 
+func failLog(err error) {
+	if err == nil {
+		return
+	}
+	log.Fatal("Failed on setup", "error", err)
+}
+
+func (c *Config) Cleanup(db *sql.DB) {
+	if c.CleanUpOnStartup != 1 {
+		return
+	}
+
+	_, err := db.Exec("update models set status=$1 where status=$2", FAILED_PREPARING_ZIP_FILE, PREPARING_ZIP_FILE)
+	failLog(err)
+	_, err = db.Exec("update models set status=$1 where status=$2", FAILED_PREPARING, PREPARING)
+	failLog(err)
+	_, err = db.Exec("update tasks set status=$1 where status=$2", TASK_PICKED_UP, TASK_TODO)
+	failLog(err)
+
+	tasks, err := GetDbMultitple[Task](db, "tasks where status=$1", TASK_RUNNING)
+	failLog(err)
+
+	base := BasePackStruct{Db: db, Logger: log.Default()}
+
+	for i := range tasks {
+		if tasks[i].TaskType == int(TASK_TYPE_CLASSIFICATION) {
+			tasks[i].UpdateStatus(base, TASK_TODO, "Reseting task")
+			continue
+		}
+		if tasks[i].TaskType == int(TASK_TYPE_RETRAINING) {
+			tasks[i].UpdateStatus(base, TASK_FAILED_RUNNING, "Task inturupted by server restart please try again")
+			_, err = db.Exec("update models set status=$1 where id=$2", READY_RETRAIN_FAILED, tasks[i].ModelId)
+			failLog(err)
+			continue
+		}
+		if tasks[i].TaskType == int(TASK_TYPE_TRAINING) {
+			tasks[i].UpdateStatus(base, TASK_FAILED_RUNNING, "Task inturupted by server restart please try again")
+			_, err = db.Exec("update models set status=$1 where id=$2", FAILED_TRAINING, tasks[i].ModelId)
+			failLog(err)
+			continue
+		}
+	}
+}
+
 func (c *Config) GenerateToken(db *sql.DB) {
 	if c.ServiceUser.User == "" {
 		log.Fatal("A user needs to be set in a configuration file")
diff --git a/main.go b/main.go
index a1c24e5..047b436 100644
--- a/main.go
+++ b/main.go
@@ -7,7 +7,6 @@ import (
 	"github.com/charmbracelet/log"
 	_ "github.com/lib/pq"
 
-	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"
 	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/models"
 	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/tasks"
 	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/tasks/runner"
@@ -15,7 +14,6 @@ import (
 )
 
 const (
-	clear_db = false
 	host     = "localhost"
 	port     = 5432
 	user     = "postgres"
@@ -46,12 +44,7 @@ func main() {
 	//TODO check if file structure exists to save data
 	handle := NewHandler(db, config)
 
-	// TODO remove this before commiting
-	_, err = db.Exec("update models set status=$1 where status=$2", FAILED_TRAINING, TRAINING)
-	if err != nil && clear_db {
-		log.Warn("Database might not be on")
-		panic(err)
-	}
+	config.Cleanup(db)
 
 	// TODO Handle this in other way
 	handle.ReadTypesFilesApi("/savedData/", ".", []string{".png", ".jpeg"}, []string{"image/png", "image/jpeg"})