Transaction RoundTripper

transaction.go

package transaction

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"sync"
)

type Transactioner interface {
	StartTransaction()
	EndTransaction()
}

type StubTransactioner struct {
	StartTransactionFn func()
	EndTransactionFn   func()
}

func (t *StubTransactioner) StartTransaction() {
	t.StartTransactionFn()
}

func (t *StubTransactioner) EndTransaction() {
	t.EndTransactionFn()
}

func newNopTransactioner() Transactioner {
	return &StubTransactioner{
		StartTransactionFn: func() {},
		EndTransactionFn:   func() {},
	}
}

type LastResponser interface {
	LastResponse() *http.Response
}

type StubLastResponser struct {
	LastResponseFn func() *http.Response
}

func (r *StubLastResponser) LastResponse() *http.Response {
	return r.LastResponseFn()
}

// TransactionalRoundTripper lets you constrain it to a transaction
// so that you can access the last response it has seen before allowing
// a different caller to use it
type TransactionalRoundTripper struct {
	rt http.RoundTripper

	mu           *sync.Mutex
	lastResponse *http.Response
}

// NewTransactionalRoundTripper creates a TransactionalRoundTripper
func NewTransactionalRoundTripper(rt http.RoundTripper) *TransactionalRoundTripper {
	return &TransactionalRoundTripper{
		rt: rt,
		mu: &sync.Mutex{},
	}
}

// RoundTrip calls the underlying RoundTripper as well as persist the response body
// so that it can be examined by the caller
func (rt *TransactionalRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := rt.rt.RoundTrip(req)
	if err != nil {
		return nil, err
	}

	// persist the response body
	respBody, err := NewRepeatingReader(resp.Body)
	if err != nil {
		return nil, err
	}

	if err := resp.Body.Close(); err != nil {
		return nil, err
	}

	// Overwrite body
	resp.Body = ioutil.NopCloser(respBody)

	// set the last response
	rt.lastResponse = resp

	return resp, nil
}

// StartTransaction starts a transaction for the round-tripper
func (rt *TransactionalRoundTripper) StartTransaction() {
	rt.mu.Lock()
}

// EndTransaction ends a transaction for the round-tripper
func (rt *TransactionalRoundTripper) EndTransaction() {
	rt.mu.Unlock()
}

// LastResponse returns the last response seen by the round-tripper
func (rt *TransactionalRoundTripper) LastResponse() *http.Response {
	return rt.lastResponse
}

// RepeatingReader is used to wrap a reader so that it can be read multiple times
type RepeatingReader struct {
	buf *bytes.Reader
}

// NewRepeatingReader creates a new RepeatingReader
func NewRepeatingReader(r io.Reader) (io.Reader, error) {
	bs, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	return &RepeatingReader{
		buf: bytes.NewReader(bs),
	}, nil
}

// Read reads from the underlying bytes reader and seeks back to
// its beginning when we are done reading
func (r *RepeatingReader) Read(b []byte) (int, error) {
	n, err := r.buf.Read(b)

	if err == io.EOF {
		if _, err := r.buf.Seek(0, io.SeekStart); err != nil {
			return 0, err
		}
	}

	return n, err
}