-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
229 lines (201 loc) · 5.8 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
package main
import (
"context"
"flag"
"fmt"
"log"
"net"
"net/http"
"net/http/httputil"
"net/url"
"strings"
"sync"
"sync/atomic"
"time"
)
// Backend provides the necessary information need for a Backend
type Backend struct {
URL *url.URL
Alive bool
mux sync.RWMutex
ReverseProxy *httputil.ReverseProxy
}
// ServerPool keeps track of available Backends
type ServerPool struct {
backends []*Backend
current uint64
}
const (
// Attempts keeps tracks of the attempts made in a Backend
Attempts int = iota
// Retry keeps tracks of the retries make in a Backend
Retry
)
// AddBackend to the server pool
func (s *ServerPool) AddBackend(backend *Backend) {
s.backends = append(s.backends, backend)
}
// NextIndex is responsible for atomicaly getting the next backend index
func (s *ServerPool) NextIndex() int {
return int(atomic.AddUint64(&s.current, uint64(1)) % uint64(len(s.backends)))
}
// MarkBackendStatus changes a status of a backend
func (s *ServerPool) MarkBackendStatus(backendURL *url.URL, alive bool) {
for _, b := range s.backends {
if b.URL.String() == backendURL.String() {
b.SetAlive(alive)
break
}
}
}
// GetNextPeer returns the next available and alive backend to take a connection
func (s *ServerPool) GetNextPeer() *Backend {
// Loop entire backends list to find an Alive backend
next := s.NextIndex()
l := len(s.backends) + next // start from next and move a full cycle
for i := next; i < l; i++ {
idx := i % len(s.backends) // take an index by modding with length of backends
// if we have an alive backend, use it and store if its not the original one
if s.backends[idx].IsAlive() {
if i != next {
atomic.StoreUint64(&s.current, uint64(idx)) // mark the current one
}
return s.backends[idx]
}
}
return nil
}
// SetAlive safely modifies the alive property
func (b *Backend) SetAlive(alive bool) {
b.mux.Lock()
b.Alive = alive
b.mux.Unlock()
}
// IsAlive safely return the alive property
func (b *Backend) IsAlive() (alive bool) {
b.mux.RLock()
alive = b.Alive
b.mux.RUnlock()
return
}
// loadBalance balances the incoming requests and only fails in
// case all backends are offline
func loadBalance(w http.ResponseWriter, r *http.Request) {
attempts := GetAttemptsFromContext(r)
if attempts > 3 {
log.Printf("%s(%s) Max attempts reached, terminating.\n", r.RemoteAddr, r.URL.Path)
http.Error(w, "Service not availabble.", http.StatusServiceUnavailable)
return
}
peer := serverPool.GetNextPeer()
if peer != nil {
peer.ReverseProxy.ServeHTTP(w, r)
return
}
http.Error(w, "Service not available", http.StatusServiceUnavailable)
}
// GetAttemptsFromContext returns the attempts for request
func GetAttemptsFromContext(r *http.Request) int {
if attempts, ok := r.Context().Value(Attempts).(int); ok {
return attempts
}
return 1
}
// GetRetryFromContext returns the retries for request
func GetRetryFromContext(r *http.Request) int {
if retry, ok := r.Context().Value(Retry).(int); ok {
return retry
}
return 0
}
// isBackendAlive checks wheter a backend is Alive by establishing a TCP connection
func isBackendAlive(u *url.URL) bool {
timeout := 2 * time.Second
conn, err := net.DialTimeout("tcp", u.Host, timeout)
if err != nil {
log.Printf("Site unreachable, error: ", err)
return false
}
_ = conn.Close() // Close the connection, we don't need to maintain it
return true
}
// HealthCheck pings the backends and update the status
func (s *ServerPool) HealthCheck() {
for _, b := range s.backends {
status := "up"
alive := isBackendAlive(b.URL)
b.SetAlive(alive)
if !alive {
status = "down"
}
log.Printf("%s [%s]\n", b.URL, status)
}
}
// healthCheck runs a routine for checking the status of all backends every 20 seconds
func healthCheck() {
t := time.NewTicker(time.Second * 20)
for {
select {
case <-t.C:
log.Printf("Starting health check...")
serverPool.HealthCheck()
log.Println("Health check finished.")
}
}
}
var serverPool ServerPool
func main() {
var serverList string
var port int
flag.StringVar(&serverList, "backends", "", "Load balanced backends, use commas to separate")
flag.IntVar(&port, "port", 3030, "Port to serve")
flag.Parse()
if len(serverList) == 0 {
log.Fatal("Please provide one or more backends to load balance")
}
// parse servers
tokens := strings.Split(serverList, ",")
for _, tok := range tokens {
serverURL, err := url.Parse(tok)
if err != nil {
log.Fatal(err)
}
reverseProxy := httputil.NewSingleHostReverseProxy(serverURL)
reverseProxy.ErrorHandler = func(writer http.ResponseWriter, request *http.Request, e error) {
log.Printf("[%s] %s\n", serverURL.Host, e.Error())
retries := GetRetryFromContext(request)
if retries < 3 {
select {
case <-time.After(10 * time.Millisecond):
ctx := context.WithValue(request.Context(), Retry, retries+1)
reverseProxy.ServeHTTP(writer, request.WithContext(ctx))
}
return
}
// after 3 retries, mark this backend as down
serverPool.MarkBackendStatus(serverURL, false)
// if the same request routing for few attempts with different backends, increase the count
attempts := GetAttemptsFromContext(request)
log.Printf("%s(%s) Attempting retry %d\n", request.RemoteAddr, request.URL.Path, attempts)
ctx := context.WithValue(request.Context(), Attempts, attempts+1)
loadBalance(writer, request.WithContext(ctx))
}
serverPool.AddBackend(&Backend{
URL: serverURL,
Alive: true,
ReverseProxy: reverseProxy,
})
log.Printf("Configured server: %s\n", serverURL)
}
// create http server
server := http.Server{
Addr: fmt.Sprintf(":%d", port),
Handler: http.HandlerFunc(loadBalance),
}
// start health checking
go healthCheck()
log.Printf("Load Balancer started at :%d\n", port)
if err := server.ListenAndServe(); err != nil {
log.Fatal(err)
}
}