Skip to content

Commit

Permalink
Tune timeouts, reorganize a little
Browse files Browse the repository at this point in the history
  • Loading branch information
xeome committed Jul 25, 2023
1 parent fc90841 commit ac052c2
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 39 deletions.
17 changes: 12 additions & 5 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package main

import (
"flag"
"log"
"os"
"time"

Expand All @@ -14,6 +13,7 @@ import (
)

var logger service.Logger
var log = utilities.Log

type program struct {
exit chan struct{}
Expand All @@ -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 {
Expand All @@ -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.")
}
}()

Expand Down
73 changes: 39 additions & 34 deletions utilities/commandline.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,52 +18,31 @@ type session struct {
var (
shellQueue *list.List
taskQueue chan struct{}
wg *sync.WaitGroup
Wg *sync.WaitGroup
)

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)

Expand Down Expand Up @@ -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:
Expand All @@ -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)
}
}

0 comments on commit ac052c2

Please sign in to comment.