CoderCookE/goaround

package pool

import (
    "bufio"
    "bytes"
    "context"
    "fmt"
    "io/ioutil"
    "log"
    "math"
    "math/rand"
    "net"
    "net/http"
    "net/http/httputil"
    "net/url"
    "os"
    "strings"
    "sync"
    "time"

    "github.com/dgraph-io/ristretto"

    "github.com/CoderCookE/goaround/internal/connection"
    "github.com/CoderCookE/goaround/internal/healthcheck"
    "github.com/CoderCookE/goaround/internal/stats"
)

type attempts int

const (
    attemptsKey attempts = iota
)

type pool struct {
    sync.RWMutex
    connections     chan *connection.Connection
    healthChecks    map[string]*healthcheck.HealthChecker
    client          *http.Client
    connsPerBackend int
    cache           *ristretto.Cache
    maxRetries      int
}

//Exported method for creation of a connection-pool takes []string
//ex: ['http://localhost:9000','http://localhost:9000']
func New(c *Config) *pool {
    backends := c.Backends
    connsPerBackend := c.NumConns
    cacheEnabled := c.EnableCache
    maxRetries := c.MaxRetries

    backendCount := int(math.Max(float64(len(backends)), float64(1)))
    maxRequests := connsPerBackend * backendCount * 2

    tr := &http.Transport{
        DialContext: (&net.Dialer{
            Timeout:   10 * time.Second,
            KeepAlive: 10 * time.Second,
        }).DialContext,
        TLSHandshakeTimeout:   10 * time.Second,
        ExpectContinueTimeout: 10 * time.Second,
        ResponseHeaderTimeout: 10 * time.Second,
    }

    client := &http.Client{
        Timeout:   30 * time.Second,
        Transport: tr,
    }

    cache, err := buildCache(cacheEnabled)
    if err != nil {
        log.Printf("Error creating cache: %v", err)
    }

    connectionPool := &pool{
        connections:     make(chan *connection.Connection, maxRequests),
        healthChecks:    make(map[string]*healthcheck.HealthChecker),
        client:          client,
        connsPerBackend: connsPerBackend,
        cache:           cache,
        maxRetries:      maxRetries,
    }

    poolConnections := []*connection.Connection{}

    startup := &sync.WaitGroup{}
    for _, backend := range backends {
        startup.Add(1)
        poolConnections = connectionPool.addBackend(poolConnections, backend, startup)
    }

    shuffle(poolConnections, connectionPool.connections)

    go connectionPool.ListenForBackendChanges(startup)

    return connectionPool
}

func shuffle(conns []*connection.Connection, ch chan *connection.Connection) {
    rand.Seed(time.Now().UnixNano())
    rand.Shuffle(len(conns), func(i, j int) {
        conns[i], conns[j] = conns[j], conns[i]
    })

    for _, conn := range conns {
        stats.AvailableConnectionsGauge.WithLabelValues("available").Add(1)
        ch <- conn
    }
}

func buildCache(cacheEnabled bool) (*ristretto.Cache, error) {
    if !cacheEnabled {
        return nil, nil
    }

    cache, err := ristretto.NewCache(&ristretto.Config{
        NumCounters: 1e7,     // number of keys to track frequency of (10M).
        MaxCost:     1 << 30, // maximum cost of cache (1GB).
        BufferItems: 64,      // number of keys per Get buffer.
    })

    return cache, err
}

//Exported method for passing a request to a connection from the pool
//Returns a 503 status code if request is unsuccessful
func (p *pool) Fetch(w http.ResponseWriter, r *http.Request) {
    start := time.Now()
    conn := <-p.connections

    ctxAttempt := r.Context().Value(attemptsKey)
    var attempt int

    if ctxAttempt != nil {
        attempt = ctxAttempt.(int) + 1
    }

    if attempt > p.maxRetries {
        return
    }

    duration := time.Since(start).Seconds()
    stats.Durations.WithLabelValues("get_connection").Observe(duration)
    stats.AvailableConnectionsGauge.WithLabelValues("in_use").Add(1)
    defer func() {
        stats.AvailableConnectionsGauge.WithLabelValues("in_use").Sub(1)
        stats.Attempts.WithLabelValues().Observe(float64(attempt))
        duration = time.Since(start).Seconds()
        stats.Durations.WithLabelValues("return_connection").Observe(duration)

        if !conn.Shut {
            p.connections <- conn
        }
    }()

    if p.cache != nil && r.Method == "GET" {
        value, found := p.cache.Get(r.URL.Path)
        if found {
            stats.CacheCounter.WithLabelValues(r.URL.Path, "hit").Add(1)
            res := value.(string)
            _, err := w.Write([]byte(res))
            if err != nil {
                log.Printf("Error writing: %s", err.Error())
            }
            return
        }

        stats.CacheCounter.WithLabelValues(r.URL.Path, "miss").Add(1)
    }

    usableProxy, err := conn.Get()
    ctx := context.WithValue(r.Context(), attemptsKey, attempt)

    if err != nil {
        log.Printf("retrying err with request: %s", err.Error())
        p.Fetch(w, r.WithContext(ctx))
    } else {
        usableProxy.ServeHTTP(w, r)
    }
}

func (p *pool) Shutdown() {
    p.RLock()
    for _, hc := range p.healthChecks {
        hc.Shutdown()
    }
    p.RUnlock()
}

func (p *pool) ListenForBackendChanges(startup *sync.WaitGroup) {
    const SockAddr = "/tmp/goaround.sock"

    if err := os.RemoveAll(SockAddr); err != nil {
        log.Fatal(err)
    }

    l, err := net.Listen("unix", SockAddr)
    if err != nil {
        log.Fatal("listen error:", err)
    }
    defer l.Close()

    for {
        conn, err := l.Accept()
        if err != nil {
            log.Fatal("accept error:", err)
        }

        scanner := bufio.NewScanner(conn)
        for scanner.Scan() {
            updated := strings.Split(scanner.Text(), ",")
            p.Lock()
            var currentBackends []string
            for k := range p.healthChecks {
                currentBackends = append(currentBackends, k)
            }

            added, removed := difference(currentBackends, updated)
            log.Printf("Adding: %s", added)
            log.Printf("Removing: %s", removed)

            for _, removedBackend := range removed {
                if len(added) > 0 {
                    var new string
                    new, added = added[0], added[1:]
                    endpoint, err := url.ParseRequestURI(new)
                    if err != nil {
                        log.Printf("Error adding backend, %s", new)
                    } else {
                        proxy := httputil.NewSingleHostReverseProxy(endpoint)
                        proxy.Transport = p.client.Transport
                        proxy.ErrorHandler = p.errorHandler
                        p.setupCache(proxy)

                        newHC := p.healthChecks[removedBackend].Reuse(new, proxy)
                        p.healthChecks[new] = newHC
                    }
                } else {
                    p.healthChecks[removedBackend].Shutdown()
                }

                delete(p.healthChecks, removedBackend)
            }

            poolConnections := []*connection.Connection{}

            wg := &sync.WaitGroup{}
            for _, addedBackend := range added {
                wg.Add(1)
                poolConnections = p.addBackend(poolConnections, addedBackend, wg)
            }

            shuffle(poolConnections, p.connections)
            p.Unlock()
        }
    }
}

func difference(original []string, updated []string) (added []string, removed []string) {
    oldBackends := make(map[string]bool)
    for _, i := range original {
        oldBackends[i] = true
    }

    newBackends := make(map[string]bool)
    for _, i := range updated {
        newBackends[i] = true
    }

    for _, i := range updated {
        if _, ok := oldBackends[i]; !ok {
            added = append(added, i)
        }
    }

    for _, i := range original {
        if _, ok := newBackends[i]; !ok {
            removed = append(removed, i)
        }
    }

    return
}

func (p *pool) addBackend(connections []*connection.Connection, backend string, startup *sync.WaitGroup) []*connection.Connection {
    endpoint, err := url.ParseRequestURI(backend)
    if err != nil {
        log.Printf("error parsing backend url: %s", backend)
    } else {
        proxy := httputil.NewSingleHostReverseProxy(endpoint)
        proxy.ErrorHandler = p.errorHandler
        proxy.Transport = p.client.Transport
        p.setupCache(proxy)

        backendConnections := make([]chan connection.Message, p.connsPerBackend)
        for i := 0; i < p.connsPerBackend; i++ {
            startup.Add(1)
            configuredConn := connection.NewConnection(proxy, backend, startup)
            connections = append(connections, configuredConn)
            backendConnections[i] = configuredConn.Messages
        }

        hc := healthcheck.New(
            p.client,
            backendConnections,
            backend,
            false,
        )

        p.healthChecks[backend] = hc

        go hc.Start(startup)
    }

    startup.Wait()
    return connections
}

func (p *pool) errorHandler(w http.ResponseWriter, r *http.Request, e error) {
    host := fmt.Sprintf("%s:%s", r.URL.Hostname(), r.URL.Port())

    stats.RequestCounter.WithLabelValues(host, "backend_error").Add(1)

    p.Fetch(w, r)
}

func (p *pool) setupCache(proxy *httputil.ReverseProxy) {
    if p.cache != nil {
        cacheResponse := func(r *http.Response) error {
            body, err := ioutil.ReadAll(r.Body)
            cacheable := string(body)
            r.Body = ioutil.NopCloser(bytes.NewBuffer(body))

            path := r.Request.URL.Path
            if err == nil {
                p.cache.Set(path, cacheable, 1)
            }

            return nil
        }

        proxy.ModifyResponse = cacheResponse
    }
}