adding flags for amount of servers and implementing weighted round robin balancing method

Author: SamSyntax

Makefile

@@ -1,2 +1,5 @@
-build:
+run:
 	go build -o ./lb *.go && ./lb
+build:
+	go build -o ./lb *.go 
+

README.md

@@ -72,10 +72,10 @@ Run the code using:
 ```bash
 go build *.go -o ./lb && ./lb
 ```
-OR
+or
 
 ```bash
-make build
+make run 
 ```
 The load balancer will listen on port 7000. Open your browser or use curl to send a request to localhost:7000:
 
@@ -84,6 +84,10 @@ curl http://localhost:7000
 ```
 The request will be forwarded to one of the backend servers, and you will see a response indicating the server's port.
 
+We can also specify amount of servers and balancing method (Round Robin and Weighted Round Robin) by building program and passing flags:
+```bash
+make build && ./lb -amount 5 -method rr
+```
 Example Output
 ```bash Serving requests at localhost:7000
 Spawning server: Server 0 at localhost:8000
@@ -98,7 +102,7 @@ forwarding to "localhost:8004"
 forwarding to "localhost:8000"
 ```
 Customization
-Number of Servers: You can modify the number of servers spawned by changing the argument passed to the Spawner function in main.go.
+Number of Servers: You can modify the number of servers spawned by changing the <b>-amount</b> flag value.
 Ports: The backend servers listen on ports 8000 and higher. You can modify the port range in the Spawner function.
 Dependencies
 This project does not require any third-party dependencies. It only uses the Go standard library.

loadbalancer

@@ -15,9 +15,11 @@ type ServerInterface interface {
 }
 
 type LbServer struct {
-	addr  string
-	proxy *httputil.ReverseProxy
-  name string
+	addr    string
+	proxy   *httputil.ReverseProxy
+	name    string
+	weight  int
+	current int
 }
 
 func (s LbServer) Address() string {
@@ -32,7 +34,7 @@ func (s LbServer) Serve(w http.ResponseWriter, r *http.Request) {
 	s.proxy.ServeHTTP(w, r)
 }
 
-func NewLbServer(addr string) *LbServer {
+func NewLbServer(addr string, weight int) *LbServer {
 	serverUrl, err := url.Parse(addr)
 	if err != nil {
 		fmt.Printf("Failed to parse url: %v\n", err)
@@ -41,24 +43,50 @@ func NewLbServer(addr string) *LbServer {
 	return &LbServer{
 		addr:  addr,
 		proxy: httputil.NewSingleHostReverseProxy(serverUrl),
+    weight: weight,
+    current: 0,
 	}
 }
 
 type LoadBalancer struct {
 	port            string
 	roundRobinCount int
 	servers         []LbServer
+	weidhted        bool
 }
 
-func NewLoadBalancer(port string, servers []LbServer) *LoadBalancer {
+func NewLoadBalancer(port string, servers []LbServer, weighted bool) *LoadBalancer {
 	return &LoadBalancer{
 		port:            port,
 		roundRobinCount: 0,
 		servers:         servers,
+		weidhted:        weighted,
 	}
 }
 
 func (lb *LoadBalancer) GetNextAvailableServer() LbServer {
+	if lb.weidhted {
+		return lb.getWeightedServer()
+	}
+	return lb.getRoundRobinServer()
+}
+func (lb *LoadBalancer) getWeightedServer() LbServer {
+
+	totalServers := len(lb.servers)
+	for i := 0; i < totalServers; i++ {
+		server := &lb.servers[lb.roundRobinCount%totalServers]
+		if server.current < server.weight {
+			server.current++
+			return *server
+		}
+		server.current = 0
+		lb.roundRobinCount++
+	}
+	lb.roundRobinCount++
+	return lb.servers[lb.roundRobinCount%totalServers]
+}
+func (lb *LoadBalancer) getRoundRobinServer() LbServer {
+
 	server := lb.servers[lb.roundRobinCount%len(lb.servers)]
 	for !server.IsAlive() {
 		lb.roundRobinCount++

loadbalancer.go

@@ -1,13 +1,27 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"net/http"
+	"os"
 )
 
 func main() {
-	servers := Spawner(5)
-	lb := NewLoadBalancer("7000", servers)
+	nFlag := flag.Int("amount", 1234, "Enter amount of local servers to be spawned")
+	method := flag.String("method", "rr", "Load balancing method: 'rr - Round Robin | wrr - Weighted Round Robin")
+	flag.Parse()
+	servers := Spawner(*nFlag)
+	var lb *LoadBalancer
+	switch *method {
+	case "wrr":
+		lb = NewLoadBalancer("7000", servers, true)
+	case "rr":
+		lb = NewLoadBalancer("7000", servers, false)
+	default:
+		fmt.Println("Invalid method. Use 'rr' or 'wrr'.")
+		os.Exit(1)
+	}
 	handleRedirect := func(w http.ResponseWriter, r *http.Request) {
 		lb.ServeProxy(w, r)
 	}

main.go

@@ -7,8 +7,8 @@ import (
 	"strconv"
 )
 
-func Server(port, name string) LbServer {
-	srv := NewLbServer("http://localhost" + port)
+func Server(port, name string, weight int) LbServer {
+	srv := NewLbServer("http://localhost"+port, weight)
 	srv.name = name
 
 	mux := http.NewServeMux()
@@ -34,10 +34,16 @@ func Server(port, name string) LbServer {
 }
 func Spawner(amt int) []LbServer {
 	servers := make([]LbServer, 0, amt)
+	weights := []int{5, 2, 3}
 	for i := 0; i < amt; i++ {
+		k := 0
+		if i > len(weights) {
+			k = 0
+		}
+		k++
 		name := fmt.Sprintf("Server %v", i)
 		port := ":" + strconv.Itoa(8000+i)
-		srv := Server(port, name)
+		srv := Server(port, name, weights[k])
 		servers = append(servers, srv)
 	}
 	return servers