From f3b4ff17b1d2d3c7cffb8e7fdd1b3f72ef74ec0b Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Tue, 17 Nov 2015 16:29:14 +0200
Subject: [PATCH] Added accessors for determining ServeConn call start time and
 the sequence number of request served in the connection

---
 server.go      | 28 +++++++++++++++++++++++++---
 server_test.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/server.go b/server.go
index 41784d9..90a68e2 100644
--- a/server.go
+++ b/server.go
@@ -197,6 +197,9 @@ type RequestCtx struct {
 
 	id uint64
 
+	serveConnRequestNum uint64
+	serveConnTime       time.Time
+
 	time time.Time
 
 	logger ctxLogger
@@ -269,6 +272,18 @@ func (ctx *RequestCtx) Time() time.Time {
 	return ctx.time
 }
 
+// ServeConnTime returns the time server starts serving the connection
+// the current request came from.
+func (ctx *RequestCtx) ServeConnTime() time.Time {
+	return ctx.serveConnTime
+}
+
+// ServeConnRequestNum returns request sequence number
+// for the current connection.
+func (ctx *RequestCtx) ServeConnRequestNum() uint64 {
+	return ctx.serveConnRequestNum
+}
+
 // SetConnectionClose sets 'Connection: close' response header and closes
 // connection after the RequestHandler returns.
 func (ctx *RequestCtx) SetConnectionClose() {
@@ -641,19 +656,22 @@ func (s *Server) serveConn(c net.Conn) error {
 	readTimeout := s.ReadTimeout
 	writeTimeout := s.WriteTimeout
 
+	currentTime := time.Now()
+
 	ctx := s.acquireCtx(c)
+	ctx.serveConnRequestNum = 0
+	ctx.serveConnTime = currentTime
 	var br *bufio.Reader
 	var bw *bufio.Writer
 	var dt time.Duration
 	var prevReadTime time.Time
 
 	var err error
-	var currentTime time.Time
 	var connectionClose bool
 	var errMsg string
 	for {
-		currentTime = time.Now()
 		ctx.id++
+		ctx.serveConnRequestNum++
 		ctx.time = currentTime
 
 		if readTimeout > 0 {
@@ -724,6 +742,8 @@ func (s *Server) serveConn(c net.Conn) error {
 				break
 			}
 		}
+
+		currentTime = time.Now()
 	}
 
 	if br != nil {
@@ -869,7 +889,9 @@ func (ctx *RequestCtx) Init(req *Request, remoteAddr net.Addr, logger Logger) {
 	ctx.initID()
 	req.CopyTo(&ctx.Request)
 	ctx.Response.Clear()
-	ctx.time = time.Now()
+	ctx.serveConnRequestNum = 0
+	ctx.serveConnTime = time.Now()
+	ctx.time = ctx.serveConnTime
 }
 
 var fakeServer Server
diff --git a/server_test.go b/server_test.go
index c4c4f88..8c78ef8 100644
--- a/server_test.go
+++ b/server_test.go
@@ -166,6 +166,51 @@ func TestServerConnectionClose(t *testing.T) {
 	}
 }
 
+func TestServerRequestNumAndTime(t *testing.T) {
+	n := uint64(0)
+	var connT time.Time
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			n++
+			if ctx.ServeConnRequestNum() != n {
+				t.Fatalf("unexpected request number: %d. Expecting %d", ctx.ServeConnRequestNum(), n)
+			}
+			if connT.IsZero() {
+				connT = ctx.ServeConnTime()
+			}
+			if ctx.ServeConnTime() != connT {
+				t.Fatalf("unexpected serve conn time: %s. Expecting %s", ctx.ServeConnTime(), connT)
+			}
+		},
+	}
+
+	rw := &readWriter{}
+	rw.r.WriteString("GET /foo1 HTTP/1.1\r\nHost: google.com\r\n\r\n")
+	rw.r.WriteString("GET /bar HTTP/1.1\r\nHost: google.com\r\n\r\n")
+	rw.r.WriteString("GET /baz HTTP/1.1\r\nHost: google.com\r\n\r\n")
+
+	ch := make(chan error)
+	go func() {
+		ch <- s.ServeConn(rw)
+	}()
+
+	select {
+	case err := <-ch:
+		if err != nil {
+			t.Fatalf("Unexpected error from serveConn: %s", err)
+		}
+	case <-time.After(100 * time.Millisecond):
+		t.Fatalf("timeout")
+	}
+
+	if n != 3 {
+		t.Fatalf("unexpected number of requests served: %d. Expecting %d", n, 3)
+	}
+
+	br := bufio.NewReader(&rw.w)
+	verifyResponse(t, br, 200, string(defaultContentType), "")
+}
+
 func TestServerEmptyResponse(t *testing.T) {
 	s := &Server{
 		Handler: func(ctx *RequestCtx) {