load balancer implementation, code refactoring

server/internal/balancer.go

@@ -0,0 +1,34 @@
+package internal
+
+import (
+	"container/heap"
+)
+
+type LoadBalancer struct {
+	pool Pool
+	done chan *Worker
+}
+
+func (b *LoadBalancer) Balance(work chan Process) {
+	for {
+		select {
+		case req := <-work:
+			b.dispatch(req)
+		case w := <-b.done:
+			b.completed(w)
+		}
+	}
+}
+
+func (b *LoadBalancer) dispatch(req Process) {
+	w := heap.Pop(&b.pool).(*Worker)
+	w.requests <- req
+	w.pending++
+	heap.Push(&b.pool, w)
+}
+
+func (b *LoadBalancer) completed(w *Worker) {
+	w.pending--
+	heap.Remove(&b.pool, w.index)
+	heap.Push(&b.pool, w)
+}

server/internal/memory_db.go

@@ -23,14 +23,17 @@ func (m *MemoryDB) Get(id string) (*Process, error) {
 	if !ok {
 		return nil, errors.New("no process found for the given key")
 	}
+
 	return entry.(*Process), nil
 }
 
 // Store a pointer of a process and return its id
 func (m *MemoryDB) Set(process *Process) string {
 	id := uuid.NewString()
+
 	m.table.Store(id, process)
 	process.Id = id
+
 	return id
 }
 
@@ -40,17 +43,20 @@ func (m *MemoryDB) Delete(id string) {
 }
 
 func (m *MemoryDB) Keys() *[]string {
-	running := []string{}
+	var running []string
+
 	m.table.Range(func(key, value any) bool {
 		running = append(running, key.(string))
 		return true
 	})
+
 	return &running
 }
 
 // Returns a slice of all currently stored processes progess
 func (m *MemoryDB) All() *[]ProcessResponse {
 	running := []ProcessResponse{}
+
 	m.table.Range(func(key, value any) bool {
 		running = append(running, ProcessResponse{
 			Id:       key.(string),
@@ -61,6 +67,7 @@ func (m *MemoryDB) All() *[]ProcessResponse {
 		})
 		return true
 	})
+
 	return &running
 }
 
@@ -75,12 +82,9 @@ func (m *MemoryDB) Persist() error {
 		return errors.Join(errors.New("failed to persist session"), err)
 	}
 
-	session := Session{
-		Processes: *running,
-	}
+	session := Session{Processes: *running}
 
-	err = gob.NewEncoder(fd).Encode(session)
-	if err != nil {
+	if err := gob.NewEncoder(fd).Encode(session); err != nil {
 		return errors.Join(errors.New("failed to persist session"), err)
 	}
 
@@ -113,7 +117,7 @@ func (m *MemoryDB) Restore(mq *MessageQueue, logger *slog.Logger) {
 
 		m.table.Store(proc.Id, restored)
 
-		if restored.Progress.Percentage != "-1" {
+		if restored.Progress.Status != StatusCompleted {
 			mq.Publish(restored)
 		}
 	}

server/internal/message_queue.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"errors"
 	"log/slog"
 
 	evbus "github.com/asaskevich/EventBus"
@@ -21,18 +22,18 @@ type MessageQueue struct {
 // By default it will be created with a size equals to nthe number of logical
 // CPU cores -1.
 // The queue size can be set via the qs flag.
-func NewMessageQueue(l *slog.Logger) *MessageQueue {
+func NewMessageQueue(l *slog.Logger) (*MessageQueue, error) {
 	qs := config.Instance().QueueSize
 
 	if qs <= 0 {
-		panic("invalid queue size")
+		return nil, errors.New("invalid queue size")
 	}
 
 	return &MessageQueue{
 		concurrency: qs,
 		eventBus:    evbus.New(),
 		logger:      l,
-	}
+	}, nil
 }
 
 // Publish a message to the queue and set the task to a peding state.

server/internal/pool.go

@@ -0,0 +1,16 @@
+package internal
+
+type Pool []*Worker
+
+func (h Pool) Len() int           { return len(h) }
+func (h Pool) Less(i, j int) bool { return h[i].index < h[j].index }
+func (h Pool) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+func (h *Pool) Push(x any)        { *h = append(*h, x.(*Worker)) }
+
+func (h *Pool) Pop() any {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}

server/internal/process.go

@@ -125,6 +125,8 @@ func (p *Process) Start() {
 
 	p.proc = cmd.Process
 
+	go p.SetMetadata()
+
 	// --------------- progress block --------------- //
 	var (
 		sourceChan = make(chan []byte)
@@ -139,7 +141,9 @@ func (p *Process) Start() {
 		defer func() {
 			r.Close()
 			p.Complete()
+
 			doneChan <- struct{}{}
+
 			close(sourceChan)
 			close(doneChan)
 		}()
@@ -215,6 +219,7 @@ func (p *Process) Kill() error {
 }
 
 // Returns the available format for this URL
+// TODO: Move out from process.go
 func (p *Process) GetFormatsSync() (DownloadFormats, error) {
 	cmd := exec.Command(config.Instance().DownloaderPath, p.Url, "-J")
 
@@ -356,9 +361,7 @@ func (p *Process) SetMetadata() error {
 	return nil
 }
 
-func (p *Process) getShortId() string {
-	return strings.Split(p.Id, "-")[0]
-}
+func (p *Process) getShortId() string { return strings.Split(p.Id, "-")[0] }
 
 func buildFilename(o *DownloadOutput) {
 	if o.Filename != "" && strings.Contains(o.Filename, ".%(ext)s") {

server/internal/worker.go

@@ -0,0 +1,15 @@
+package internal
+
+type Worker struct {
+	requests chan Process // downloads to do
+	pending  int          // downloads pending
+	index    int          // index in the heap
+}
+
+func (w *Worker) Work(done chan *Worker) {
+	for {
+		req := <-w.requests
+		req.Start()
+		done <- w
+	}
+}