Covered by tests

2026-09-30 23:42:31 +00:00 · 2014-12-11 15:24:48 -08:00
parent f99d039025
commit 596a7a3102
5 changed files with 254 additions and 125 deletions
@@ -1,12 +1,16 @@
-Verse
+Oxy
 =====

-Vere is a Go library for reverse proxying.
+Oxy is a Go library with HTTP handlers that enhance HTTP standard library:

-It is designed to be fully compatible with http standard library, easy to customize and re-use.
+* Stream middleware retries and buffers requests and responses http://godoc.org/github.com/mailgun/oxy/stream
+* Forward middleware forwards requests to remote location and rewrites headers http://godoc.org/github.com/mailgun/oxy/stream
+* Roundrobin middleware is a round-robin load balancer http://godoc.org/github.com/mailgun/oxy/stream

+It is designed to be fully compatible with http standard library, easy to customize and reuse.

 Status
 ------

-Just started :-)
+* Initial design is completed
+* Covered by tests
@@ -1,82 +0,0 @@
-package retry
-
-import (
-	"io"
-	"net/http"
-
-	"github.com/mailgun/oxy/utils"
-)
-
-type Retry struct {
-	next        http.Handler
-	errHandler  http.Handler
-	maxAttempts int
-	predicate   hpredicate
-}
-
-func New(next http.Handler, predicate string, settings ...optSetter) (*Retry, error) {
-	p, err := parseExpression(predicate)
-	if err != nil {
-		return nil, err
-	}
-
-	r := &Retry{
-		next:      next,
-		predicate: p,
-	}
-	for _, s := range settings {
-		if err := s(r); err != nil {
-			return nil, err
-		}
-	}
-	if r.maxAttempts == 0 {
-		r.maxAttempts = DefaultMaxAttempts
-	}
-	return r, nil
-}
-
-func (r *Retry) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	c := &context{r: req}
-	for i := 0; i < r.maxAttempts; i++ {
-		pw := &utils.ProxyWriter{W: w}
-
-		// this is not the first attempt, we need to rewind the body
-		if i != 0 {
-			seeker, ok := req.Body.(io.Seeker)
-			if !ok {
-				r.errHandler.ServeHTTP(w, req)
-				return
-			}
-			_, err := seeker.Seek(0, 0)
-			if err != nil {
-				r.errHandler.ServeHTTP(w, req)
-				return
-			}
-		}
-		r.next.ServeHTTP(pw, req)
-		c.attempt = i + 1
-		if !r.predicate(c) {
-			return
-		}
-	}
-	// if we ended up here, we have excceded all the max attempts what is a system error
-	r.errHandler.ServeHTTP(w, req)
-}
-
-func MaxAttempts(a int) optSetter {
-	return func(r *Retry) error {
-		r.maxAttempts = a
-		return nil
-	}
-}
-
-func ErrorHandler(h http.Handler) optSetter {
-	return func(r *Retry) error {
-		r.errHandler = h
-		return nil
-	}
-}
-
-type optSetter func(r *Retry) error
-
-const DefaultMaxAttempts = 10
@@ -0,0 +1,95 @@
+package stream
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"os"
+
+	"github.com/mailgun/oxy/forward"
+	"github.com/mailgun/oxy/roundrobin"
+	"github.com/mailgun/oxy/testutils"
+	"github.com/mailgun/oxy/utils"
+
+	. "gopkg.in/check.v1"
+)
+
+type RTSuite struct{}
+
+var _ = Suite(&RTSuite{})
+
+func (s *RTSuite) TestSuccess(c *C) {
+	srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
+		w.Write([]byte("hello"))
+	})
+	defer srv.Close()
+
+	lb, rt := new(c, `IsNetworkError() && Attempts() <= 2`)
+
+	proxy := httptest.NewServer(rt)
+	defer proxy.Close()
+
+	lb.UpsertServer(testutils.ParseURI(srv.URL))
+
+	re, body, err := testutils.Get(proxy.URL)
+	c.Assert(err, IsNil)
+	c.Assert(re.StatusCode, Equals, http.StatusOK)
+	c.Assert(string(body), Equals, "hello")
+}
+
+func (s *RTSuite) TestRetryOnError(c *C) {
+	srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
+		w.Write([]byte("hello"))
+	})
+	defer srv.Close()
+
+	lb, rt := new(c, `IsNetworkError() && Attempts() <= 2`)
+
+	proxy := httptest.NewServer(rt)
+	defer proxy.Close()
+
+	lb.UpsertServer(testutils.ParseURI("http://localhost:64321"))
+	lb.UpsertServer(testutils.ParseURI(srv.URL))
+
+	re, body, err := testutils.Get(proxy.URL)
+	c.Assert(err, IsNil)
+	c.Assert(re.StatusCode, Equals, http.StatusOK)
+	c.Assert(string(body), Equals, "hello")
+}
+
+func (s *RTSuite) TestRetryExceedAttempts(c *C) {
+	srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
+		w.Write([]byte("hello"))
+	})
+	defer srv.Close()
+
+	lb, rt := new(c, `IsNetworkError() && Attempts() <= 2`)
+
+	proxy := httptest.NewServer(rt)
+	defer proxy.Close()
+
+	lb.UpsertServer(testutils.ParseURI("http://localhost:64321"))
+	lb.UpsertServer(testutils.ParseURI("http://localhost:64322"))
+	lb.UpsertServer(testutils.ParseURI("http://localhost:64323"))
+	lb.UpsertServer(testutils.ParseURI(srv.URL))
+
+	re, _, err := testutils.Get(proxy.URL)
+	c.Assert(err, IsNil)
+	c.Assert(re.StatusCode, Equals, http.StatusBadGateway)
+}
+
+func new(c *C, p string) (*roundrobin.RoundRobin, *Streamer) {
+	logger := utils.NewFileLogger(os.Stdout, utils.INFO)
+	// forwarder will proxy the request to whatever destination
+	fwd, err := forward.New(forward.Logger(logger))
+	c.Assert(err, IsNil)
+
+	// load balancer will round robin request
+	lb, err := roundrobin.New(fwd)
+	c.Assert(err, IsNil)
+
+	// stream handler will forward requests to redirect, make sure it uses files
+	st, err := New(lb, Logger(logger), Retry(p), MemRequestBodyBytes(1))
+	c.Assert(err, IsNil)
+
+	return lb, st
+}
@@ -22,6 +22,8 @@ const (
 	DefaultMemBodyBytes = 1048576
 	// No limit by default
 	DefaultMaxBodyBytes = -1
+	// Maximum retry attempts
+	DefaultMaxRetryAttempts = 10
 )

 var errHandler utils.ErrorHandler = &SizeErrHandler{}
@@ -35,6 +37,8 @@ type Streamer struct {
 	maxResponseBodyBytes int64
 	memResponseBodyBytes int64

+	retryPredicate hpredicate
+
 	next       http.Handler
 	errHandler utils.ErrorHandler
 	log        utils.Logger
@@ -73,6 +77,17 @@ func New(next http.Handler, setters ...optSetter) (*Streamer, error) {

 type optSetter func(s *Streamer) error

+func Retry(predicate string) optSetter {
+	return func(s *Streamer) error {
+		p, err := parseExpression(predicate)
+		if err != nil {
+			return err
+		}
+		s.retryPredicate = p
+		return nil
+	}
+}
+
 // Logger
 func Logger(l utils.Logger) optSetter {
 	return func(s *Streamer) error {
@@ -148,7 +163,6 @@ func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	// to read into memory and disk. This reader returns an error if the total request size exceeds the
 	// prefefined MaxSizeBytes. This can occur if we got chunked request, in this case ContentLength would be set to -1
 	// and the reader would be unbounded bufio in the http.Server
-
 	body, err := multibuf.New(req.Body, multibuf.MaxBytes(s.maxRequestBodyBytes), multibuf.MemBytes(s.memRequestBodyBytes))
 	if err != nil || body == nil {
 		s.errHandler.ServeHTTP(w, req, err)
@@ -160,9 +174,6 @@ func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	// and we don'w want to mess with standard library
 	defer body.Close()

-	outreq := *req
-	outreq.Body = body
-
 	// We need to set ContentLength based on known request size. The incoming request may have been
 	// set without content length or using chunked TransferEncoding
 	totalSize, err := body.Size()
@@ -171,37 +182,66 @@ func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		s.errHandler.ServeHTTP(w, req, err)
 		return
 	}
-	outreq.ContentLength = totalSize
+
+	outreq := s.copyRequest(req, body, totalSize)
+
+	attempt := 1
+	for {
+		// We create a special writer that will limit the response size, buffer it to disk if necessary
+		writer, err := multibuf.NewWriterOnce(multibuf.MaxBytes(s.maxResponseBodyBytes), multibuf.MemBytes(s.memResponseBodyBytes))
+		if err != nil {
+			s.errHandler.ServeHTTP(w, req, err)
+			return
+		}
+
+		// We are mimicking http.ResponseWriter to replace writer with our special writer
+		b := &bufferWriter{
+			header: make(http.Header),
+			buffer: writer,
+		}
+		defer b.Close()
+
+		s.next.ServeHTTP(b, outreq)
+
+		s.log.Infof("next responded: %v", b.code)
+
+		reader, err := writer.Reader()
+		if err != nil {
+			s.log.Errorf("failed to read response, err %v", err)
+			s.errHandler.ServeHTTP(w, req, err)
+			return
+		}
+		defer reader.Close()
+
+		if (s.retryPredicate == nil || attempt > DefaultMaxRetryAttempts) ||
+			!s.retryPredicate(&context{r: req, attempt: attempt, responseCode: b.code, log: s.log}) {
+			copyHeaders(w.Header(), b.Header())
+			w.WriteHeader(b.code)
+			io.Copy(w, reader)
+			return
+		}
+
+		attempt += 1
+		if _, err := body.Seek(0, 0); err != nil {
+			s.log.Errorf("Failed to rewind: error: %v", err)
+			s.errHandler.ServeHTTP(w, req, err)
+			return
+		}
+		outreq = s.copyRequest(req, body, totalSize)
+		s.log.Infof("retry Request(%v %v) attempt %v", req.Method, req.URL, attempt)
+	}
+}
+
+func (s *Streamer) copyRequest(req *http.Request, body io.ReadCloser, bodySize int64) *http.Request {
+	o := *req
+	o.URL = utils.CopyURL(req.URL)
+	o.Header = make(http.Header)
+	utils.CopyHeaders(o.Header, req.Header)
+	o.ContentLength = bodySize
 	// remove TransferEncoding that could have been previously set because we have transformed the request from chunked encoding
-	outreq.TransferEncoding = []string{}
-
-	// We create a special writer that will limit the response size, buffer it to disk if necessary
-	writer, err := multibuf.NewWriterOnce(multibuf.MaxBytes(s.maxResponseBodyBytes), multibuf.MemBytes(s.memResponseBodyBytes))
-	if err != nil {
-		s.errHandler.ServeHTTP(w, req, err)
-		return
-	}
-
-	// We are mimicking http.ResponseWriter to replace writer with our special writer
-	b := &bufferWriter{
-		header: make(http.Header),
-		buffer: writer,
-	}
-	defer b.Close()
-
-	s.next.ServeHTTP(b, &outreq)
-
-	reader, err := writer.Reader()
-	if err != nil {
-		s.log.Errorf("failed to read response, err %v", err)
-		s.errHandler.ServeHTTP(w, req, err)
-		return
-	}
-	defer reader.Close()
-
-	copyHeaders(w.Header(), b.Header())
-	w.WriteHeader(b.code)
-	io.Copy(w, reader)
+	o.TransferEncoding = []string{}
+	o.Body = body
+	return &o
 }

 func (s *Streamer) checkLimit(req *http.Request) error {
@@ -1,10 +1,10 @@
-package retry
+package stream

 import (
 	"fmt"
-	"net"
 	"net/http"

+	"github.com/mailgun/oxy/utils"
 	"github.com/mailgun/predicate"
 )

@@ -12,7 +12,7 @@ type context struct {
 	r            *http.Request
 	attempt      int
 	responseCode int
-	err          error
+	log          utils.Logger
 }

 type hpredicate func(*context) bool
@@ -25,6 +25,10 @@ func parseExpression(in string) (hpredicate, error) {
 			OR:  or,
 			EQ:  eq,
 			NEQ: neq,
+			LT:  lt,
+			GT:  gt,
+			LE:  le,
+			GE:  ge,
 		},
 		Functions: map[string]interface{}{
 			"RequestMethod":  requestMethod,
@@ -74,8 +78,8 @@ func responseCode() toInt {
 // IsNetworkError returns a predicate that returns true if last attempt ended with network error.
 func isNetworkError() hpredicate {
 	return func(c *context) bool {
-		_, ok := c.err.(net.Error)
-		return ok
+		c.log.Infof("code: %v", c.responseCode)
+		return c.responseCode == http.StatusBadGateway || c.responseCode == http.StatusGatewayTimeout
 	}
 }

@@ -130,6 +134,54 @@ func neq(m interface{}, value interface{}) (hpredicate, error) {
 	return not(p), nil
 }

+// lt returns predicate that tests that value of the mapper function is less than the constant
+func lt(m interface{}, value interface{}) (hpredicate, error) {
+	switch mapper := m.(type) {
+	case toInt:
+		return intLT(mapper, value)
+	}
+	return nil, fmt.Errorf("unsupported argument: %T", m)
+}
+
+// le returns predicate that tests that value of the mapper function is less or equal than the constant
+func le(m interface{}, value interface{}) (hpredicate, error) {
+	l, err := lt(m, value)
+	if err != nil {
+		return nil, err
+	}
+	e, err := eq(m, value)
+	if err != nil {
+		return nil, err
+	}
+	return func(c *context) bool {
+		return l(c) || e(c)
+	}, nil
+}
+
+// gt returns predicate that tests that value of the mapper function is greater than the constant
+func gt(m interface{}, value interface{}) (hpredicate, error) {
+	switch mapper := m.(type) {
+	case toInt:
+		return intGT(mapper, value)
+	}
+	return nil, fmt.Errorf("unsupported argument: %T", m)
+}
+
+// ge returns predicate that tests that value of the mapper function is less or equal than the constant
+func ge(m interface{}, value interface{}) (hpredicate, error) {
+	g, err := gt(m, value)
+	if err != nil {
+		return nil, err
+	}
+	e, err := eq(m, value)
+	if err != nil {
+		return nil, err
+	}
+	return func(c *context) bool {
+		return g(c) || e(c)
+	}, nil
+}
+
 func stringEQ(m toString, val interface{}) (hpredicate, error) {
 	value, ok := val.(string)
 	if !ok {
@@ -149,3 +201,23 @@ func intEQ(m toInt, val interface{}) (hpredicate, error) {
 		return m(c) == value
 	}, nil
 }
+
+func intLT(m toInt, val interface{}) (hpredicate, error) {
+	value, ok := val.(int)
+	if !ok {
+		return nil, fmt.Errorf("expected int, got %T", val)
+	}
+	return func(c *context) bool {
+		return m(c) < value
+	}, nil
+}
+
+func intGT(m toInt, val interface{}) (hpredicate, error) {
+	value, ok := val.(int)
+	if !ok {
+		return nil, fmt.Errorf("expected int, got %T", val)
+	}
+	return func(c *context) bool {
+		return m(c) > value
+	}, nil
+}