-
Notifications
You must be signed in to change notification settings - Fork 565
/
balancing.go
78 lines (67 loc) · 2.45 KB
/
balancing.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
package proxy
import (
"context"
"net/url"
"strings"
"github.com/devopsfaith/krakend/config"
"github.com/devopsfaith/krakend/sd"
)
// NewLoadBalancedMiddleware creates proxy middleware adding the most perfomant balancer
// over a default subscriber
func NewLoadBalancedMiddleware(remote *config.Backend) Middleware {
return NewLoadBalancedMiddlewareWithSubscriber(sd.GetSubscriber(remote))
}
// NewLoadBalancedMiddlewareWithSubscriber creates proxy middleware adding the most perfomant balancer
// over the received subscriber
func NewLoadBalancedMiddlewareWithSubscriber(subscriber sd.Subscriber) Middleware {
return newLoadBalancedMiddleware(sd.NewBalancer(subscriber))
}
// NewRoundRobinLoadBalancedMiddleware creates proxy middleware adding a round robin balancer
// over a default subscriber
func NewRoundRobinLoadBalancedMiddleware(remote *config.Backend) Middleware {
return NewRoundRobinLoadBalancedMiddlewareWithSubscriber(sd.GetSubscriber(remote))
}
// NewRandomLoadBalancedMiddleware creates proxy middleware adding a random balancer
// over a default subscriber
func NewRandomLoadBalancedMiddleware(remote *config.Backend) Middleware {
return NewRandomLoadBalancedMiddlewareWithSubscriber(sd.GetSubscriber(remote))
}
// NewRoundRobinLoadBalancedMiddlewareWithSubscriber creates proxy middleware adding a round robin
// balancer over the received subscriber
func NewRoundRobinLoadBalancedMiddlewareWithSubscriber(subscriber sd.Subscriber) Middleware {
return newLoadBalancedMiddleware(sd.NewRoundRobinLB(subscriber))
}
// NewRandomLoadBalancedMiddlewareWithSubscriber creates proxy middleware adding a random
// balancer over the received subscriber
func NewRandomLoadBalancedMiddlewareWithSubscriber(subscriber sd.Subscriber) Middleware {
return newLoadBalancedMiddleware(sd.NewRandomLB(subscriber))
}
func newLoadBalancedMiddleware(lb sd.Balancer) Middleware {
return func(next ...Proxy) Proxy {
if len(next) > 1 {
panic(ErrTooManyProxies)
}
return func(ctx context.Context, request *Request) (*Response, error) {
host, err := lb.Host()
if err != nil {
return nil, err
}
r := request.Clone()
var b strings.Builder
b.WriteString(host)
b.WriteString(r.Path)
r.URL, err = url.Parse(b.String())
if err != nil {
return nil, err
}
if len(r.Query) > 0 {
if len(r.URL.RawQuery) > 0 {
r.URL.RawQuery += "&" + r.Query.Encode()
} else {
r.URL.RawQuery += r.Query.Encode()
}
}
return next[0](ctx, &r)
}
}
}