From 95f266f1ceeab1c9007ea9bc63670d3afa2d04a3 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Thu, 22 Feb 2024 09:42:12 +0100
Subject: [PATCH] tsweb: add optional on completion callback func

Updates corp#17075

Co-Authored-By: Anton Tolchanov <anton@tailscale.com>
Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 tsweb/tsweb.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index d3bdaefe5..f1ba6d965 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -254,11 +254,19 @@ type HandlerOptions struct {
 	// is intended to be used to present pretty error pages if
 	// the user agent is determined to be a browser.
 	OnError ErrorHandlerFunc
+
+	// OnCompletion is called when ServeHTTP is finished and gets
+	// useful data that the implementor can use for metrics.
+	OnCompletion OnCompletionFunc
 }
 
 // ErrorHandlerFunc is called to present a error response.
 type ErrorHandlerFunc func(http.ResponseWriter, *http.Request, HTTPError)
 
+// OnCompletionFunc is called when ServeHTTP is finished and gets
+// useful data that the implementor can use for metrics.
+type OnCompletionFunc func(*http.Request, AccessLogRecord)
+
 // ReturnHandlerFunc is an adapter to allow the use of ordinary
 // functions as ReturnHandlers. If f is a function with the
 // appropriate signature, ReturnHandlerFunc(f) is a ReturnHandler that
@@ -435,6 +443,10 @@ func (h retHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
+	if h.opts.OnCompletion != nil {
+		h.opts.OnCompletion(r, msg)
+	}
+
 	if bs := h.opts.BucketedStats; bs != nil && bs.Finished != nil {
 		// Only increment metrics for buckets that result in good HTTP statuses
 		// or when we know the start was already counted.