openbilibili/app/service/live/broadcast-proxy/server/proxy.go

package server

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"go-common/library/log"
)

type BroadcastProxy struct {
	backend         []string
	probePath       string
	reverseProxy    []*httputil.ReverseProxy
	bestClientIndex int32
	probeSample     int
	wg              sync.WaitGroup
	ctx             context.Context
	cancel          context.CancelFunc
}

func NewBroadcastProxy(backend []string, probePath string, maxIdleConns int, probeSample int) (*BroadcastProxy, error) {
	if len(backend) == 0 {
		return nil, errors.New("Require at least one backend")
	}
	proxy := new(BroadcastProxy)
	if probeSample > 0 {
		proxy.probeSample = probeSample
	} else {
		proxy.probeSample = 1
	}
	for _, addr := range backend {
		proxy.backend = append(proxy.backend, addr)
		proxy.probePath = probePath
		p := httputil.NewSingleHostReverseProxy(&url.URL{
			Scheme: "http",
			Host:   addr,
		})
		p.Transport = &http.Transport{
			DisableKeepAlives:   false,
			MaxIdleConns:        maxIdleConns,
			MaxIdleConnsPerHost: maxIdleConns,
		}
		proxy.reverseProxy = append(proxy.reverseProxy, p)
	}
	proxy.ctx, proxy.cancel = context.WithCancel(context.Background())
	proxy.wg.Add(1)
	go func() {
		defer proxy.wg.Done()
		proxy.mainProbeProcess()
	}()

	return proxy, nil
}

func (proxy *BroadcastProxy) Close() {
	proxy.cancel()
	proxy.wg.Wait()
}

func (proxy *BroadcastProxy) HandleRequest(w http.ResponseWriter, r *http.Request) {
	t1 := time.Now()
	i := atomic.LoadInt32(&proxy.bestClientIndex)
	proxy.reverseProxy[i].ServeHTTP(w, r)
	t2 := time.Now()
	log.V(3).Info("proxy process req:%s,backend id:%d, timecost:%s", r.RequestURI, i, t2.Sub(t1).String())
}

func (proxy *BroadcastProxy) RequestAllBackend(method, uri string, requestBody []byte) ([]string, []error) {
	responseCollection := make([]string, len(proxy.reverseProxy))
	errCollection := make([]error, len(proxy.reverseProxy))
	var wg sync.WaitGroup
	for i := range proxy.reverseProxy {
		wg.Add(1)
		go func(index int) {
			defer wg.Done()
			req, err := http.NewRequest(method, uri, bytes.NewReader(requestBody))
			if err != nil {
				errCollection[index] = err
				return
			}
			httpRecorder := httptest.NewRecorder()
			proxy.reverseProxy[index].ServeHTTP(httpRecorder, req)
			resp := httpRecorder.Result()
			defer resp.Body.Close()
			if resp.StatusCode != http.StatusOK {
				errCollection[index] = errors.New(fmt.Sprintf("http response:%s", resp.Status))
				return
			}
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				errCollection[index] = err
				return
			}
			responseCollection[index] = string(body)
		}(i)
	}
	wg.Wait()
	return responseCollection, errCollection
}

func (proxy *BroadcastProxy) mainProbeProcess() {
	var wg sync.WaitGroup
	interval := make([]time.Duration, len(proxy.backend))
	for {
		fast := -1
		for i, probe := range proxy.reverseProxy {
			wg.Add(1)
			go func(p *httputil.ReverseProxy, index int) {
				defer wg.Done()
				interval[index] = proxy.unitProbeProcess(p, index)
			}(probe, i)
		}
		wg.Wait()
		for i, v := range interval {
			if fast < 0 {
				fast = i
			} else {
				if v < interval[fast] {
					fast = i
				}
			}
		}
		atomic.StoreInt32(&proxy.bestClientIndex, int32(fast))
		log.Info("[probe result]best server id:%d,addr:%s", fast, proxy.backend[fast])
		for i, d := range interval {
			log.Info("[probe log]server id:%d,addr:%s,avg time cost:%fms", i, proxy.backend[i], 1000*d.Seconds())
		}
		select {
		case <-time.After(time.Second):
		case <-proxy.ctx.Done():
			return
		}
	}
}

func (proxy *BroadcastProxy) unitProbeProcess(p *httputil.ReverseProxy, backendIndex int) time.Duration {
	var (
		wg       sync.WaitGroup
		duration int64
	)
	for i := 0; i < proxy.probeSample; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			timeout := time.Second
			<-time.After(time.Duration(rand.Intn(proxy.probeSample)) * time.Millisecond)
			if timeCost, err := proxy.checkMonitorPing(p, backendIndex, timeout); err == nil {
				atomic.AddInt64(&duration, timeCost.Nanoseconds())
			} else {
				atomic.AddInt64(&duration, timeout.Nanoseconds())
			}
		}()
	}
	wg.Wait()
	return time.Duration(duration/int64(proxy.probeSample)) * time.Nanosecond
}

func (proxy *BroadcastProxy) checkMonitorPing(p *httputil.ReverseProxy, backendIndex int, timeout time.Duration) (time.Duration, error) {
	req, err := http.NewRequest("GET", proxy.probePath, nil)
	if err != nil {
		return timeout, err
	}
	ctx, cancel := context.WithTimeout(proxy.ctx, timeout)
	defer cancel()
	req = req.WithContext(ctx)

	recorder := httptest.NewRecorder()
	beginTime := time.Now()
	p.ServeHTTP(recorder, req)
	resp := recorder.Result()
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		log.Error("probe:server id:%d,addr:%s,send requset error:%s", backendIndex,
			proxy.backend[backendIndex], resp.Status)
		return timeout, errors.New("http response:" + resp.Status)
	}
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Error("probe:server id:%d,addr:%s,read response error:%v", backendIndex,
			proxy.backend[backendIndex], err)
		return timeout, err
	}
	if !bytes.Equal(result, []byte("pong")) {
		log.Error("probe:server id:%d,addr:%s,not match response:%s", backendIndex,
			proxy.backend[backendIndex], result)
		return timeout, err
	}
	endTime := time.Now()
	return endTime.Sub(beginTime), nil
}