From f79c689c5da7e037a048b5e1fa4ca2e06879aaf2 Mon Sep 17 00:00:00 2001
From: Paul Kulchenko <paul@kulchenko.com>
Date: Sun, 17 Jul 2022 22:10:24 -0700
Subject: [PATCH] Add redbean OnLogLatency hook.

---
 tool/net/help.txt  |  6 ++++++
 tool/net/redbean.c | 31 ++++++++++++++++++++++++++-----
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/tool/net/help.txt b/tool/net/help.txt
index 2d63ec00d..96219f1c3 100644
--- a/tool/net/help.txt
+++ b/tool/net/help.txt
@@ -534,6 +534,12 @@ HOOKS
           each time redbean accepts a new client connection. If it returns
           `true`, redbean will close the connection without calling fork.
 
+  OnLogLatency(reqtimeus:int,contimeus:int)
+          If this function is defined it'll be called from the main process
+          each time redbean completes handling of a request, but before the
+          response is sent. The handler received the time (in µs) since the
+          request handling and connection handling started.
+
   OnProcessCreate(pid:int,ip:int,port:int,serverip:int,serverport:int)
           If this function is defined it'll be called from the main process
           each time redbean forks a connection handler worker process. The
diff --git a/tool/net/redbean.c b/tool/net/redbean.c
index a4527d921..b77852d84 100644
--- a/tool/net/redbean.c
+++ b/tool/net/redbean.c
@@ -385,6 +385,7 @@ static bool gotcachecontrol;
 static bool interpretermode;
 static bool sslclientverify;
 static bool connectionclose;
+static bool hasonloglatency;
 static bool hasonworkerstop;
 static bool isexitingworker;
 static bool hasonworkerstart;
@@ -1108,6 +1109,21 @@ static bool LuaOnClientConnection(void) {
   return dropit;
 }
 
+static void LuaOnLogLatency(long reqtime, long contime) {
+#ifndef STATIC
+  lua_State *L = GL;
+  int n = lua_gettop(L);
+  lua_getglobal(L, "OnLogLatency");
+  lua_pushinteger(L, reqtime);
+  lua_pushinteger(L, contime);
+  if (LuaCallWithTrace(L, 2, 0, NULL) != LUA_OK) {
+    LogLuaError("OnLogLatency", lua_tostring(L, -1));
+    lua_pop(L, 1);  // pop error
+  }
+  AssertLuaStackIsAt(L, n);
+#endif
+}
+
 static void LuaOnProcessCreate(int pid) {
 #ifndef STATIC
   uint32_t ip, serverip;
@@ -5483,6 +5499,7 @@ static void LuaInit(void) {
     hasonprocessdestroy = IsHookDefined("OnProcessDestroy");
     hasonworkerstart = IsHookDefined("OnWorkerStart");
     hasonworkerstop = IsHookDefined("OnWorkerStop");
+    hasonloglatency = IsHookDefined("OnLogLatency");
   } else {
     DEBUGF("(srvr) no /.init.lua defined");
   }
@@ -6258,6 +6275,7 @@ static bool StreamResponse(char *p) {
 
 static bool HandleMessageActual(void) {
   int rc;
+  long reqtime, contime;
   char *p;
   long double now;
   if ((rc = ParseHttpMessage(&msg, inbuf.p, amtread)) != -1) {
@@ -6298,12 +6316,15 @@ static bool HandleMessageActual(void) {
     p = stpcpy(p, referrerpolicy);
     p = stpcpy(p, "\r\n");
   }
-  if (loglatency || LOGGABLE(kLogDebug)) {
+  if (loglatency || LOGGABLE(kLogDebug) || hasonloglatency) {
     now = nowl();
-    LOGF(kLogDebug, "(stat) %`'.*s latency r: %,ldµs c: %,ldµs",
-         msg.uri.b - msg.uri.a, inbuf.p + msg.uri.a,
-         (long)((now - startrequest) * 1e6L),
-         (long)((now - startconnection) * 1e6L));
+    reqtime = (long)((now - startrequest) * 1e6L);
+    contime = (long)((now - startconnection) * 1e6L);
+    if (hasonloglatency) LuaOnLogLatency(reqtime, contime);
+    if (loglatency || LOGGABLE(kLogDebug))
+      LOGF(kLogDebug, "(stat) %`'.*s latency r: %,ldµs c: %,ldµs",
+           msg.uri.b - msg.uri.a, inbuf.p + msg.uri.a,
+           reqtime, contime);
   }
   if (!generator) {
     return TransmitResponse(p);