From a8a3770e8db025abcd13f86f8bb6c260a409fdb6 Mon Sep 17 00:00:00 2001
From: Ajay kumar dhyani <dhyani.devops@gmail.com>
Date: Wed, 30 Oct 2024 00:58:06 -0700
Subject: [PATCH] feat: improve code

---
 example/simple.go | 15 +++++++++++++++
 loadbalancer.go   | 39 ++++++++++++++++++++++++---------------
 server.go         | 25 +++++++++++++++++--------
 3 files changed, 56 insertions(+), 23 deletions(-)
 create mode 100644 example/simple.go

diff --git a/example/simple.go b/example/simple.go
new file mode 100644
index 0000000..896bc37
--- /dev/null
+++ b/example/simple.go
@@ -0,0 +1,15 @@
+package main
+
+import (
+	"github.com/dhyanio/golb"
+)
+
+func main() {
+	serverList := []string{
+		"https://www.google.com/",
+		"https://github.com/dhyanio",
+		"https://www.linkedin.com/in/dhyanio/",
+	}
+
+	golb.Start(serverList, "3000")
+}
diff --git a/loadbalancer.go b/loadbalancer.go
index e52afc1..a2ee567 100644
--- a/loadbalancer.go
+++ b/loadbalancer.go
@@ -1,16 +1,20 @@
 package golb
 
 import (
-	"fmt"
+	"log"
 	"net/http"
+	"sync"
 )
 
+// LoadBalancer represents a simple round-robin load balancer.
 type LoadBalancer struct {
 	port            string
 	roundRobinCount int
 	servers         []Server
+	mu              sync.Mutex
 }
 
+// NewLoadBalancer initializes a LoadBalancer with the specified port and servers.
 func NewLoadBalancer(port string, servers []Server) *LoadBalancer {
 	return &LoadBalancer{
 		port:            port,
@@ -18,7 +22,12 @@ func NewLoadBalancer(port string, servers []Server) *LoadBalancer {
 		servers:         servers,
 	}
 }
+
+// getNextAvailableServer returns the next available server in a round-robin fashion.
 func (lb *LoadBalancer) getNextAvailableServer() Server {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
 	server := lb.servers[lb.roundRobinCount%len(lb.servers)]
 	for !server.IsAlive() {
 		lb.roundRobinCount++
@@ -28,26 +37,26 @@ func (lb *LoadBalancer) getNextAvailableServer() Server {
 	return server
 }
 
-func (lb *LoadBalancer) serverProxy(rw http.ResponseWriter, r *http.Request) {
-	targetSever := lb.getNextAvailableServer()
-	fmt.Printf("Forwarding request to address %q\n", targetSever.Address())
-	targetSever.Serve(rw, r)
+// serveProxy forwards the request to the selected server.
+func (lb *LoadBalancer) serveProxy(rw http.ResponseWriter, r *http.Request) {
+	targetServer := lb.getNextAvailableServer()
+	log.Printf("Forwarding request to address %q\n", targetServer.Address())
+	targetServer.Serve(rw, r)
 }
 
-// Start starts loadbalancer
-func start(serverAdd []string, port string) {
+// Start initializes servers, sets up the load balancer, and starts listening on the specified port.
+func Start(serverAddresses []string, port string) {
 	var servers []Server
 
-	for _, value := range serverAdd {
-		servers = append(servers, newSimpleServer(value))
+	for _, addr := range serverAddresses {
+		servers = append(servers, NewSimpleServer(addr))
 	}
 
 	lb := NewLoadBalancer(port, servers)
-	handleRedirect := func(rw http.ResponseWriter, req *http.Request) {
-		lb.serverProxy(rw, req)
-	}
+	http.HandleFunc("/", lb.serveProxy)
 
-	http.HandleFunc("/", handleRedirect)
-	fmt.Printf("Serving requests at localhost:%s\n", lb.port)
-	http.ListenAndServe(":"+lb.port, nil)
+	log.Printf("Load balancer listening on port %s\n", port)
+	if err := http.ListenAndServe(":"+port, nil); err != nil {
+		log.Fatalf("Failed to start server: %v", err)
+	}
 }
diff --git a/server.go b/server.go
index c6f6908..3fea128 100644
--- a/server.go
+++ b/server.go
@@ -6,32 +6,41 @@ import (
 	"net/url"
 )
 
-// Server represents server
+// Server represents a load-balanced server interface.
 type Server interface {
 	Address() string
 	IsAlive() bool
-	Serve(rw http.ResponseWriter, r *http.Request)
+	Serve(http.ResponseWriter, *http.Request)
 }
 
-type simpleServer struct {
+// SimpleServer represents a basic HTTP server with reverse proxy capabilities.
+type SimpleServer struct {
 	addr  string
 	proxy *httputil.ReverseProxy
 }
 
-func newSimpleServer(addr string) *simpleServer {
+// NewSimpleServer creates a new instance of SimpleServer with a reverse proxy.
+func NewSimpleServer(addr string) *SimpleServer {
 	serverUrl, err := url.Parse(addr)
 	handleErr(err)
 
-	return &simpleServer{
+	return &SimpleServer{
 		addr:  addr,
 		proxy: httputil.NewSingleHostReverseProxy(serverUrl),
 	}
 }
 
-func (s *simpleServer) Address() string { return s.addr }
+// Address returns the server address.
+func (s *SimpleServer) Address() string {
+	return s.addr
+}
 
-func (s *simpleServer) IsAlive() bool { return true }
+// IsAlive always returns true, but this could be expanded for real health checks.
+func (s *SimpleServer) IsAlive() bool {
+	return true
+}
 
-func (s *simpleServer) Serve(rw http.ResponseWriter, req *http.Request) {
+// Serve handles the HTTP request by forwarding it to the reverse proxy.
+func (s *SimpleServer) Serve(rw http.ResponseWriter, req *http.Request) {
 	s.proxy.ServeHTTP(rw, req)
 }