Merge branch 'feat/logging' (pull/3)

* feat/logging:
  fix: allow ResponseReader.Reader to be wrapped by transcoders
  Add logging for startup and transcoded bytes

compy.go

@@ -6,6 +6,9 @@ import (
 	"github.com/barnacs/compy/proxy"
 	tc "github.com/barnacs/compy/transcoder"
 	"log"
+	"os"
+	"os/signal"
+	"sync/atomic"
 )
 
 var (
@@ -55,6 +58,20 @@ func main() {
 	p.AddTranscoder("application/javascript", ttc)
 	p.AddTranscoder("application/x-javascript", ttc)
 
+	c := make(chan os.Signal, 2)
+	signal.Notify(c, os.Interrupt)
+	go func() {
+		for _ = range c {
+			read := atomic.LoadUint64(&p.ReadCount)
+			written := atomic.LoadUint64(&p.WriteCount)
+			log.Printf("compy exiting, total transcoded: %d -> %d (%3.1f%%)",
+				read, written, float64(written)/float64(read)*100)
+			os.Exit(0)
+		}
+	}()
+
+	log.Printf("compy listening on %s", *host)
+
 	var err error
 	if *cert != "" {
 		err = p.StartTLS(*host, *cert, *key)

proxy/proxy.go

@@ -5,11 +5,14 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"sync/atomic"
 )
 
 type Proxy struct {
 	transcoders map[string]Transcoder
 	ml          *mitmListener
+	ReadCount   uint64
+	WriteCount  uint64
 }
 
 type Transcoder interface {
@@ -47,6 +50,7 @@ func (p *Proxy) StartTLS(host, cert, key string) error {
 }
 
 func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	log.Printf("serving request: %s", r.URL)
 	if err := p.handle(w, r); err != nil {
 		log.Printf("%s while serving request: %s", err, r.URL)
 	}
@@ -62,7 +66,15 @@ func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error {
 		return fmt.Errorf("error forwarding request: %s", err)
 	}
 	defer resp.Body.Close()
-	return p.proxyResponse(newResponseWriter(w), newResponseReader(resp))
+	rw := newResponseWriter(w)
+	rr := newResponseReader(resp)
+	err = p.proxyResponse(rw, rr)
+	read := rr.counter.Count()
+	written := rw.rw.Count()
+	log.Printf("transcoded: %d -> %d (%3.1f%%)", read, written, float64(written)/float64(read)*100)
+	atomic.AddUint64(&p.ReadCount, read)
+	atomic.AddUint64(&p.WriteCount, written)
+	return err
 }
 
 func forward(r *http.Request) (*http.Response, error) {

proxy/response.go

@@ -4,16 +4,21 @@ import (
 	"io"
 	"mime"
 	"net/http"
+
+	"github.com/miolini/datacounter"
 )
 
 type ResponseReader struct {
 	io.Reader
+	counter *datacounter.ReaderCounter
 	r       *http.Response
 }
 
 func newResponseReader(r *http.Response) *ResponseReader {
+	counter := datacounter.NewReaderCounter(r.Body)
 	return &ResponseReader{
-		Reader: r.Body,
+		Reader:  counter,
+		counter: counter,
 		r:       r,
 	}
 }
@@ -34,15 +39,16 @@ func (r *ResponseReader) Request() *http.Request {
 
 type ResponseWriter struct {
 	io.Writer
-	rw          http.ResponseWriter
+	rw          *datacounter.ResponseWriterCounter
 	statusCode  int
 	headersDone bool
 }
 
 func newResponseWriter(w http.ResponseWriter) *ResponseWriter {
+	rw := datacounter.NewResponseWriterCounter(w)
 	return &ResponseWriter{
-		Writer: w,
+		Writer: rw,
-		rw:     w,
+		rw:     rw,
 	}
 }