func (p *LoadBalancer) printError(rw http.ResponseWriter, down *DownStream, msg string) { dlog.Warn(msg) rw.WriteHeader(500) rw.Write([]byte(msg)) if down != nil { down.used = false } }
func (p *LoadBalancer) ServeHTTP(rw http.ResponseWriter, req *http.Request) { srv, ok := p.services[req.Host] if !ok { p.printError(rw, nil, "fail to get service: "+req.Host) return } down := srv.createDownstream(req) if nil == down { dlog.Warn("fail to find downstream") rw.WriteHeader(500) p.printError(rw, nil, "fail to find downstream") return } down.used = true body, err := ioutil.ReadAll(req.Body) if err != nil { p.printError(rw, down, "fail to read body: "+err.Error()) return } req2, err := http.NewRequest(req.Method, "http://"+down.up.Addr+req.RequestURI, bytes.NewReader(body)) dlog.Println(req2.Method, " ", req2.URL.String()) if err != nil { p.printError(rw, down, "new req2 fail: "+err.Error()) return } resp, err := p.client.Do(req2) if err != nil { p.printError(rw, down, "call upstream failed: "+err.Error()) return } if resp.Body == nil { p.printError(rw, down, "upstream resp body is nil") return } defer resp.Body.Close() b, err := ioutil.ReadAll(resp.Body) if err != nil { p.printError(rw, down, "read upstream body fail: "+err.Error()) return } http.SetCookie(rw, &http.Cookie{ Name: "lbc_session_id", Value: down.sid, }) rw.WriteHeader(200) rw.Write(b) down.used = false }
func NewServiceHandler(srv *ServiceConfig) *ServiceHandler { ret := &ServiceHandler{ srv: srv, freeUpstreams: make(chan *UpstreamConfig, 100), cache: make(map[string]*DownStream), ticker: time.NewTicker(time.Second * 3), } for _, up := range srv.Upstreams { dlog.Println("append up: ", up.Addr) ret.freeUpstreams <- up } go func() { for t := range ret.ticker.C { dlog.Warn("close timeout upstream: %d", t.Unix()) del := []*DownStream{} for _, d := range ret.cache { if d.used && time.Now().Sub(d.startTime).Seconds() > 300 { d.used = false del = append(del, d) continue } if !d.used && time.Now().Sub(d.startTime).Seconds() > 120 { del = append(del, d) continue } } for _, d := range del { dlog.Println("free upstream: ", d.up.Addr) delete(ret.cache, d.sid) ret.freeUpstreams <- d.up } } }() return ret }