Create stats endpoint

Dockerfile

@@ -6,7 +6,7 @@ COPY go.mod ./
 
 COPY . .
 
-RUN CGO_ENABLED=0 GOOS=linux go build -o lb cmd/lb/main.go
+RUN CGO_ENABLED=0 GOOS=linux go build -o lb ./cmd/lb
 
 FROM alpine:latest
 

Dockerfile.backend

@@ -0,0 +1,19 @@
+FROM golang:1.25.4-alpine AS builder
+
+WORKDIR /app
+
+COPY go.mod ./
+
+COPY . .
+
+RUN CGO_ENABLED=0 GOOS=linux go build -o server cmd/server/main.go
+
+FROM alpine:latest
+
+WORKDIR /root/
+
+COPY --from=builder /app/server .
+
+EXPOSE 80
+
+ENTRYPOINT ["./server"]

cmd/lb/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"flag"
 	"fmt"
 	"log"
@@ -59,10 +60,40 @@ func healthCheck() {
 
 				b.SetAlive(alive)
 			}
+
+			if alive {
+				go updateBackendStats(b)
+			}
 		}
 	}
 }
 
+type HealthResponse struct {
+	MemoryUsage uint64 `json:"memory_usage"`
+}
+
+func updateBackendStats(b *core.Backend) {
+	resp, err := http.Get(b.URL.String() + "/health")
+	if err != nil {
+		log.Printf("Error fetching stats from %s: %s", b.URL, err)
+		return
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		log.Printf("Error fetching stats from %s: status %d", b.URL, resp.StatusCode)
+		return
+	}
+
+	var health HealthResponse
+	if err := json.NewDecoder(resp.Body).Decode(&health); err != nil {
+		log.Printf("Error decoding stats from %s: %s", b.URL, err)
+		return
+	}
+
+	b.SetMemoryUsage(health.MemoryUsage)
+}
+
 // isBackendAlive checks whether a backend is alive by establishing a TCP connection
 func isBackendAlive(u *url.URL) bool {
 	timeout := 2 * time.Second
@@ -89,6 +120,10 @@ func main() {
 		log.Fatal("Please provide one or more backends using -backends")
 	}
 
+	// Register handlers
+	http.HandleFunc("/", lbHandler)
+	http.HandleFunc("/stats", statsHandler)
+
 	// Parse servers
 	tokens := strings.Split(serverList, ",")
 	for _, tok := range tokens {
@@ -112,14 +147,15 @@ func main() {
 			URL:          serverUrl,
 			Alive:        true, // We assume they are alive at startup
 			ReverseProxy: proxy,
+			StartTime:    time.Now(),
 		})
 		log.Printf("Configured server: %s\n", serverUrl)
 	}
 
 	// Create HTTP server
 	server := http.Server{
 		Addr:    fmt.Sprintf(":%d", port),
-		Handler: http.HandlerFunc(lbHandler),
+		Handler: nil, // Use DefaultServeMux
 	}
 
 	// Start health checking in a separate goroutine

cmd/lb/stats.go

@@ -0,0 +1,17 @@
+package main
+
+import (
+	"encoding/json"
+	"net/http"
+)
+
+// statsHandler returns the current status of the server pool
+func statsHandler(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+
+	stats := serverPool.GetStats()
+
+	if err := json.NewEncoder(w).Encode(stats); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+}

cmd/server/main.go

@@ -0,0 +1,40 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"runtime"
+)
+
+type HealthResponse struct {
+	MemoryUsage uint64 `json:"memory_usage"`
+}
+
+func main() {
+	port := os.Getenv("PORT")
+	if port == "" {
+		port = "80"
+	}
+
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(w, "Hello from backend! I am running on %s\n", os.Getenv("HOSTNAME"))
+	})
+
+	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+		var m runtime.MemStats
+		runtime.ReadMemStats(&m)
+
+		resp := HealthResponse{
+			MemoryUsage: m.Alloc,
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(resp)
+	})
+
+	log.Printf("Backend server starting on port %s...", port)
+	log.Fatal(http.ListenAndServe(":"+port, nil))
+}

core/backend.go

@@ -1,9 +1,11 @@
 package core
 
 import (
+	"fmt"
 	"net/http/httputil"
 	"net/url"
 	"sync"
+	"time"
 )
 
 // Backend is our unique server instance
@@ -13,13 +15,20 @@ type Backend struct {
 	Alive        bool
 	Mux          sync.RWMutex
 	ReverseProxy *httputil.ReverseProxy
+	StartTime    time.Time
+	MemoryUsage  uint64
 }
 
 // SetAlive is a thread-safe way to set the alive status of the backend
 func (b *Backend) SetAlive(alive bool) {
 	b.Mux.Lock()
+	defer b.Mux.Unlock()
+
+	// If we are transitioning to alive, reset the timer
+	if alive && !b.Alive {
+		b.StartTime = time.Now()
+	}
 	b.Alive = alive
-	b.Mux.Unlock()
 }
 
 // IsAlive is a thread-safe way to read the alive status of the backend
@@ -29,3 +38,52 @@ func (b *Backend) IsAlive() (alive bool) {
 	b.Mux.RUnlock()
 	return
 }
+
+// GetUpTime returns the uptime in a human-readable format
+func (b *Backend) GetUpTime() string {
+	return formatSecondsToDuration(b.GetUpTimeInSeconds())
+}
+
+// GetUpTimeInSeconds returns the uptime in seconds
+func (b *Backend) GetUpTimeInSeconds() uint64 {
+	b.Mux.RLock()
+	defer b.Mux.RUnlock()
+	if !b.Alive {
+		return 0
+	}
+	return uint64(time.Since(b.StartTime).Seconds())
+}
+
+// SetMemoryUsage sets the memory usage of the backend
+func (b *Backend) SetMemoryUsage(mem uint64) {
+	b.Mux.Lock()
+	defer b.Mux.Unlock()
+	b.MemoryUsage = mem
+}
+
+// GetMemoryUsage returns the current memory usage
+func (b *Backend) GetMemoryUsage() uint64 {
+	b.Mux.RLock()
+	defer b.Mux.RUnlock()
+	return b.MemoryUsage
+}
+
+// GetMemoryUsageString returns the memory usage in a human-readable format
+func (b *Backend) GetMemoryUsageString() string {
+	mem := b.GetMemoryUsage()
+	if mem < 1024 {
+		return fmt.Sprintf("%d B", mem)
+	} else if mem < 1024*1024 {
+		return fmt.Sprintf("%.2f KB", float64(mem)/1024)
+	} else {
+		return fmt.Sprintf("%.2f MB", float64(mem)/(1024*1024))
+	}
+}
+
+// BackendStats represents the statistics of a backend server
+type BackendStats struct {
+	URL         string `json:"url"`
+	Alive       bool   `json:"alive"`
+	UpTime      string `json:"uptime"`
+	MemoryUsage string `json:"memory_usage"`
+}

core/pool.go

@@ -1,6 +1,7 @@
 package core
 
 import (
+	"fmt"
 	"sync/atomic"
 )
 
@@ -31,3 +32,51 @@ func (s *ServerPool) GetNextPeer() *Backend {
 func (s *ServerPool) AddBackend(b *Backend) {
 	s.Backends = append(s.Backends, b)
 }
+
+// formatSecondsToDuration converts seconds to a human-readable duration string
+func formatSecondsToDuration(seconds uint64) string {
+	hours := seconds / 3600
+	minutes := (seconds % 3600) / 60
+	secs := seconds % 60
+	return fmt.Sprintf("%02dh:%02dm:%02ds", hours, minutes, secs)
+}
+
+// GetUpTimeInSeconds returns the total uptime in seconds of all alive backends
+func (s *ServerPool) GetUpTimeInSeconds() uint64 {
+	var total uint64
+	for _, b := range s.Backends {
+		if b.IsAlive() {
+			total += b.GetUpTimeInSeconds()
+		}
+	}
+	return total
+}
+
+func (s *ServerPool) GetUpTime() string {
+	var totalUpTime uint64
+	var aliveCount uint64
+	for _, b := range s.Backends {
+		if b.IsAlive() {
+			totalUpTime += b.GetUpTimeInSeconds()
+			aliveCount++
+		}
+	}
+	if aliveCount == 0 {
+		return "0s"
+	}
+	averageUpTime := totalUpTime / aliveCount
+	return formatSecondsToDuration(averageUpTime)
+}
+
+func (s *ServerPool) GetStats() []BackendStats {
+	var stats []BackendStats
+	for _, b := range s.Backends {
+		stats = append(stats, BackendStats{
+			URL:         b.URL.String(),
+			Alive:       b.IsAlive(),
+			UpTime:      b.GetUpTime(),
+			MemoryUsage: b.GetMemoryUsageString(),
+		})
+	}
+	return stats
+}

docker-compose.yml

@@ -13,17 +13,29 @@ services:
       - app2
       - app3
 
-  # Backend 1 (Simulated)
+  # Backend 1 (Custom Go Server)
   app1:
-    image: traefik/whoami
+    build:
+      context: .
+      dockerfile: Dockerfile.backend
     container_name: app1
+    environment:
+      - PORT=80
 
-  # Backend 2 (Simulated)
+  # Backend 2 (Custom Go Server)
   app2:
-    image: traefik/whoami
+    build:
+      context: .
+      dockerfile: Dockerfile.backend
     container_name: app2
+    environment:
+      - PORT=80
 
-  # Backend 3 (Simulated)
+  # Backend 3 (Custom Go Server)
   app3:
-    image: traefik/whoami
-    container_name: app3
+    build:
+      context: .
+      dockerfile: Dockerfile.backend
+    container_name: app3
+    environment:
+      - PORT=80