Ensure at least one log flush happens before request ends (#164)

Force a log flush before the connection closes. Fixes #163

internal/api.go

@@ -129,7 +129,13 @@ func handleHTTP(w http.ResponseWriter, r *http.Request) {
 		flushes++
 	}
 	c.pendingLogs.Unlock()
-	go c.flushLog(false)
+	flushed := make(chan struct{})
+	go func() {
+		defer close(flushed)
+		// Force a log flush, because with very short requests we
+		// may not ever flush logs.
+		c.flushLog(true)
+	}()
 	w.Header().Set(logFlushHeader, strconv.Itoa(flushes))
 
 	// Avoid nil Write call if c.Write is never called.
@@ -139,6 +145,9 @@ func handleHTTP(w http.ResponseWriter, r *http.Request) {
 	if c.outBody != nil {
 		w.Write(c.outBody)
 	}
+	// Wait for the last flush to complete before returning,
+	// otherwise the security ticket will not be valid.
+	<-flushed
 }
 
 func executeRequestSafely(c *context, r *http.Request) {

internal/api_test.go

@@ -250,43 +250,78 @@ func TestDelayedLogFlushing(t *testing.T) {
 	f, c, cleanup := setup()
 	defer cleanup()
 
-	http.HandleFunc("/quick_log", func(w http.ResponseWriter, r *http.Request) {
+	http.HandleFunc("/slow_log", func(w http.ResponseWriter, r *http.Request) {
 		logC := WithContext(netcontext.Background(), r)
 		fromContext(logC).apiURL = c.apiURL // Otherwise it will try to use the default URL.
 		Logf(logC, 1, "It's a lovely day.")
 		w.WriteHeader(200)
+		time.Sleep(1200 * time.Millisecond)
 		w.Write(make([]byte, 100<<10)) // write 100 KB to force HTTP flush
 	})
 
 	r := &http.Request{
 		Method: "GET",
 		URL: &url.URL{
 			Scheme: "http",
-			Path:   "/quick_log",
+			Path:   "/slow_log",
 		},
 		Header: c.req.Header,
 		Body:   ioutil.NopCloser(bytes.NewReader(nil)),
 	}
 	w := httptest.NewRecorder()
 
-	// Check that log flushing does not hold up the HTTP response.
-	start := time.Now()
-	handleHTTP(w, r)
-	if d := time.Since(start); d > 10*time.Millisecond {
-		t.Errorf("handleHTTP took %v, want under 10ms", d)
+	handled := make(chan struct{})
+	go func() {
+		defer close(handled)
+		handleHTTP(w, r)
+	}()
+	// Check that the log flush eventually comes in.
+	time.Sleep(1200 * time.Millisecond)
+	if f := atomic.LoadInt32(&f.LogFlushes); f != 1 {
+		t.Errorf("After 1.2s: f.LogFlushes = %d, want 1", f)
 	}
+
+	<-handled
 	const hdr = "X-AppEngine-Log-Flush-Count"
-	if h := w.HeaderMap.Get(hdr); h != "1" {
-		t.Errorf("%s header = %q, want %q", hdr, h, "1")
+	if got, want := w.HeaderMap.Get(hdr), "1"; got != want {
+		t.Errorf("%s header = %q, want %q", hdr, got, want)
 	}
-	if f := atomic.LoadInt32(&f.LogFlushes); f != 0 {
-		t.Errorf("After HTTP response: f.LogFlushes = %d, want 0", f)
+	if got, want := atomic.LoadInt32(&f.LogFlushes), int32(2); got != want {
+		t.Errorf("After HTTP response: f.LogFlushes = %d, want %d", got, want)
 	}
 
-	// Check that the log flush eventually comes in.
-	time.Sleep(100 * time.Millisecond)
-	if f := atomic.LoadInt32(&f.LogFlushes); f != 1 {
-		t.Errorf("After 100ms: f.LogFlushes = %d, want 1", f)
+}
+
+func TestLogFlushing(t *testing.T) {
+	f, c, cleanup := setup()
+	defer cleanup()
+
+	http.HandleFunc("/quick_log", func(w http.ResponseWriter, r *http.Request) {
+		logC := WithContext(netcontext.Background(), r)
+		fromContext(logC).apiURL = c.apiURL // Otherwise it will try to use the default URL.
+		Logf(logC, 1, "It's a lovely day.")
+		w.WriteHeader(200)
+		w.Write(make([]byte, 100<<10)) // write 100 KB to force HTTP flush
+	})
+
+	r := &http.Request{
+		Method: "GET",
+		URL: &url.URL{
+			Scheme: "http",
+			Path:   "/quick_log",
+		},
+		Header: c.req.Header,
+		Body:   ioutil.NopCloser(bytes.NewReader(nil)),
+	}
+	w := httptest.NewRecorder()
+
+	handleHTTP(w, r)
+	const hdr = "X-AppEngine-Log-Flush-Count"
+	if got, want := w.HeaderMap.Get(hdr), "1"; got != want {
+		t.Errorf("%s header = %q, want %q", hdr, got, want)
+	}
+	if got, want := atomic.LoadInt32(&f.LogFlushes), int32(1); got != want {
+		t.Errorf("After HTTP response: f.LogFlushes = %d, want %d", got, want)
 	}
 }