Add multiple remotes and round robin

This commit is contained in:
Pablu23
2025-05-19 09:30:07 +02:00
parent 403c89b068
commit 55a2d2b708
5 changed files with 73 additions and 36 deletions

View File

@@ -7,38 +7,48 @@ import (
"io"
"net/http"
"net/http/httputil"
"sync/atomic"
"github.com/pablu23/domain-router/util"
"github.com/rs/zerolog/log"
)
type Router struct {
config *Config
domains *util.ImmutableMap[string, Host]
domains *ThreadMap[string, Host]
client *http.Client
}
type Host struct {
Port int
Remote string
Secure bool
Port int
Remotes []string
Secure bool
Current *atomic.Uint32
}
func New(config *Config, client *http.Client) Router {
m := make(map[string]Host)
for _, host := range config.Hosts {
for _, domain := range host.Domains {
m[domain] = Host{host.Port, host.Remote, host.Secure}
m[domain] = Host{host.Port, host.Remotes, host.Secure, &atomic.Uint32{}}
}
}
return Router{
config: config,
domains: util.NewImmutableMap(m),
domains: NewThreadMap(m),
client: client,
}
}
func (router *Router) roundRobin(host *Host) {
l := len(host.Remotes)
if l > 1 && host.Current.Load()+1 < uint32(l) {
host.Current.Add(1)
} else if l > 1 {
host.Current.Store(0)
}
}
func (router *Router) Healthz(w http.ResponseWriter, r *http.Request) {
if !router.config.General.AnnouncePublic {
http.NotFound(w, r)
@@ -58,9 +68,9 @@ func (router *Router) Healthz(w http.ResponseWriter, r *http.Request) {
healthy := true
var url string
if host.Secure {
url = fmt.Sprintf("https://%s:%d/healthz", host.Remote, host.Port)
url = fmt.Sprintf("https://%s:%d/healthz", host.Remotes, host.Port)
} else {
url = fmt.Sprintf("http://%s:%d/healthz", host.Remote, host.Port)
url = fmt.Sprintf("http://%s:%d/healthz", host.Remotes, host.Port)
}
res, err := router.client.Get(url)
@@ -89,13 +99,13 @@ func (router *Router) Healthz(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}
func createRequest(r *http.Request, host *Host) (*http.Request, error) {
func createRequest(r *http.Request, host *Host, remote string) (*http.Request, error) {
subUrlPath := r.URL.RequestURI()
var url string
if host.Secure {
url = fmt.Sprintf("https://%s:%d%s", host.Remote, host.Port, subUrlPath)
url = fmt.Sprintf("https://%s:%d%s", remote, host.Port, subUrlPath)
} else {
url = fmt.Sprintf("http://%s:%d%s", host.Remote, host.Port, subUrlPath)
url = fmt.Sprintf("http://%s:%d%s", remote, host.Port, subUrlPath)
}
req, err := http.NewRequest(r.Method, url, r.Body)
@@ -165,9 +175,12 @@ func (router *Router) Route(w http.ResponseWriter, r *http.Request) {
return
}
req, err := createRequest(r, &host)
remote := host.Remotes[host.Current.Load()]
go router.roundRobin(&host)
req, err := createRequest(r, &host, remote)
if err != nil {
log.Error().Err(err).Bool("secure", host.Secure).Str("remote", host.Remote).Int("port", host.Port).Msg("Could not create request")
log.Error().Err(err).Bool("secure", host.Secure).Str("remote", remote).Int("port", host.Port).Msg("Could not create request")
w.WriteHeader(http.StatusInternalServerError)
return
}
@@ -179,7 +192,7 @@ func (router *Router) Route(w http.ResponseWriter, r *http.Request) {
res, err := router.client.Do(req)
if err != nil {
log.Error().Err(err).Str("remote", host.Remote).Int("port", host.Port).Msg("Could not complete request")
log.Error().Err(err).Str("remote", remote).Int("port", host.Port).Msg("Could not complete request")
w.WriteHeader(http.StatusInternalServerError)
return
}