From ac052c232f1cbb6f2fa834276cf8834b6360f4a1 Mon Sep 17 00:00:00 2001
From: Emin <emin.delice@mono.net.tr>
Date: Tue, 25 Jul 2023 15:45:18 +0300
Subject: [PATCH] Tune timeouts, reorganize a little

---
 main.go                  | 17 +++++++---
 utilities/commandline.go | 73 +++++++++++++++++++++-------------------
 2 files changed, 51 insertions(+), 39 deletions(-)

diff --git a/main.go b/main.go
index 87222d3..360ac43 100644
--- a/main.go
+++ b/main.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"flag"
-	"log"
 	"os"
 	"time"
 
@@ -14,6 +13,7 @@ import (
 )
 
 var logger service.Logger
+var log = utilities.Log
 
 type program struct {
 	exit chan struct{}
@@ -34,7 +34,7 @@ func (p *program) Start(s service.Service) error {
 
 func (p *program) run() {
 	c := utilities.ParseConfig()
-	rest.StartServer(c.Port, "1.3.0")
+	rest.StartServer(c.Port, "1.3.1")
 }
 
 func (p *program) Stop(s service.Service) error {
@@ -53,21 +53,28 @@ func main() {
 
 	utilities.Init()
 	utilities.SetupLogger()
+	utilities.Wg.Add(1)
+	go func() {
+		time.Sleep(10 * time.Second)
+		log.Info("Application started.")
+		utilities.Wg.Done()
+	}()
+
 	hyperv.Init()
 
 	go func() {
 		for {
 			time.Sleep(660 * time.Second)
 			hyperv.Refresh()
-			logger.Info("Hyper-V module reinitialized.")
+			log.Info("Hyper-V module reinitialized.")
 		}
 	}()
 
 	go func() {
 		for {
-			time.Sleep(2700 * time.Second)
+			time.Sleep(1919 * time.Second)
 			utilities.RefreshShellQueue()
-			logger.Info("Shell queue reinitialized.")
+			log.Info("Shell queue reinitialized.")
 		}
 	}()
 
diff --git a/utilities/commandline.go b/utilities/commandline.go
index d696e41..9db8dd4 100644
--- a/utilities/commandline.go
+++ b/utilities/commandline.go
@@ -18,7 +18,7 @@ type session struct {
 var (
 	shellQueue *list.List
 	taskQueue  chan struct{}
-	wg         *sync.WaitGroup
+	Wg         *sync.WaitGroup
 )
 
 const maxQueueSize = 5
@@ -26,44 +26,23 @@ const maxQueueSize = 5
 func Init() {
 	shellQueue = list.New()
 	taskQueue = make(chan struct{}, maxQueueSize)
-	wg = &sync.WaitGroup{}
-}
-
-func addSession() {
-	wg.Wait()
-	if shellQueue.Len() < maxQueueSize {
-		shell, err := powershell.New(&backend.Local{})
-		if err != nil {
-			panic(err)
-		}
-		newSession := &session{shell: &shell, busy: false}
-		shellQueue.PushBack(newSession)
-		// If there are sessions in the queue and taskQueue is not full, release a slot for the next task
-		if taskQueue != nil && len(taskQueue) < cap(taskQueue) {
-			taskQueue <- struct{}{}
-			return
-		}
-		log.Warn("No task queue or task queue is full")
-	}
-}
-
-func rotateQueue() {
-	if shellQueue.Len() > 0 {
-		e := shellQueue.Front()
-		shellQueue.MoveToBack(e)
-	}
+	Wg = &sync.WaitGroup{}
 }
 
 func CommandLine(command string) ([]byte, error) {
-	wg.Wait()
+	Wg.Wait()
 	for shellQueue.Len() < maxQueueSize {
-		addSession()
+		addShell()
 	}
 
 	if taskQueue == nil {
 		return nil, fmt.Errorf("task queue is not initialized")
 	}
 
+	return dispatchTask(command)
+}
+
+func dispatchTask(command string) ([]byte, error) {
 	errChan := make(chan error)
 	result := make(chan []byte)
 
@@ -95,8 +74,9 @@ func CommandLine(command string) ([]byte, error) {
 		}(s)
 
 		select {
-		case <-time.After(300 * time.Second):
+		case <-time.After(1080 * time.Second):
 			s.busy = false
+			(*s.shell).Exit()
 			go RefreshShellQueue()
 			return nil, fmt.Errorf("timeout")
 		case err := <-errChan:
@@ -108,17 +88,42 @@ func CommandLine(command string) ([]byte, error) {
 	return nil, fmt.Errorf("no session available")
 }
 
+func addShell() {
+	Wg.Wait()
+	if shellQueue.Len() < maxQueueSize {
+		shell, err := powershell.New(&backend.Local{})
+		if err != nil {
+			panic(err)
+		}
+		newSession := &session{shell: &shell, busy: false}
+		shellQueue.PushBack(newSession)
+		// If there are sessions in the queue and taskQueue is not full, release a slot for the next task
+		if taskQueue != nil && len(taskQueue) < cap(taskQueue) {
+			taskQueue <- struct{}{}
+			return
+		}
+		log.Warn("No task queue or task queue is full")
+	}
+}
+
+func rotateQueue() {
+	if shellQueue.Len() > 0 {
+		e := shellQueue.Front()
+		shellQueue.MoveToBack(e)
+	}
+}
+
 func RefreshShellQueue() {
-	wg.Add(1)
-	defer wg.Done()
-	for shellQueue.Len() > 0 {
+	Wg.Add(1)
+	defer Wg.Done()
+	for shellQueue.Len() > 1 {
 		e := shellQueue.Front()
 		s := e.Value.(*session)
 		if s.busy {
 			rotateQueue()
 			continue
 		}
-		shellQueue.Remove(e)
 		(*s.shell).Exit()
+		shellQueue.Remove(e)
 	}
 }