Add shutdown hook function for lua plugin (#13045)

* Add shutdown hook function for lua plugin

* Update lua function name

* update comments

* Fix problem with calling the lua function

* Update plugins/lua/ts_lua.cc

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix format

* Change lua function name

* Fix test formatting

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: shukitchan

doc/admin-guide/plugins/lua.en.rst

@@ -125,6 +125,17 @@ each lua script:
 - **'do_global_send_response'**
 - **'do_global_cache_lookup_complete'**
 - **'do_global_read_cache'**
+- **'__shutdown__'**
+
+The ``__shutdown__`` function is invoked once per Lua state when |ATS| is
+shutting down. It can be used to perform cleanup tasks such as flushing state or
+releasing resources. It takes no arguments and its return value is ignored.
+
+Example::
+
+    function __shutdown__()
+        ts.debug('ATS shutting down, cleaning up resources')
+    end
 
 We can write this in plugin.config:
 

plugins/lua/ts_lua.cc

@@ -827,6 +827,47 @@ globalHookHandler(TSCont contp, TSEvent event ATS_UNUSED, void *edata)
   return 0;
 }
 
+static int
+shutdownHookHandler(TSCont contp, TSEvent /* event ATS_UNUSED */, void * /* edata ATS_UNUSED */)
+{
+  ts_lua_instance_conf *const conf = (ts_lua_instance_conf *)TSContDataGet(contp);
+
+  for (int index = 0; index < conf->states; ++index) {
+    ts_lua_main_ctx *const main_ctx = &ts_lua_g_main_ctx_array[index];
+
+    TSMutexLock(main_ctx->mutexp);
+
+    lua_State *const L = main_ctx->lua;
+
+    // Restore the conf-specific global table so lua_getglobal resolves
+    // functions from the loaded script, matching ts_lua_reload_module.
+    lua_pushlightuserdata(L, conf);
+    lua_rawget(L, LUA_REGISTRYINDEX);
+    lua_replace(L, LUA_GLOBALSINDEX);
+
+    lua_getglobal(L, TS_LUA_FUNCTION_G_SHUT_DOWN);
+
+    if (lua_type(L, -1) == LUA_TFUNCTION) {
+      if (lua_pcall(L, 0, 0, 0) != 0) {
+        TSError("[ts_lua][%s] lua_pcall failed for script '%s' state %d: %s", __FUNCTION__, conf->script, index,
+                lua_tostring(L, -1));
+        lua_pop(L, 1);
+      }
+    } else {
+      lua_pop(L, 1);
+    }
+
+    // Restore LUA_GLOBALSINDEX to an empty table, matching the resting state
+    // established by ts_lua_add_module and ts_lua_reload_module.
+    lua_newtable(L);
+    lua_replace(L, LUA_GLOBALSINDEX);
+
+    TSMutexUnlock(main_ctx->mutexp);
+  }
+
+  return 0;
+}
+
 void
 TSPluginInit(int argc, const char *argv[])
 {
@@ -1046,7 +1087,7 @@ TSPluginInit(int argc, const char *argv[])
   }
   TSContDataSet(vconn_contp, conf);
 
-  // adding hook based on whther the lua global vconn function exists
+  // adding hook based on whether the lua global vconn function exists
   ts_lua_vconn_ctx *vconn_ctx = ts_lua_create_vconn_ctx(main_ctx, conf);
   lua_State        *vl        = vconn_ctx->lua;
 
@@ -1059,6 +1100,30 @@ TSPluginInit(int argc, const char *argv[])
 
   ts_lua_destroy_vconn_ctx(vconn_ctx);
 
+  // adding shutdown hook if the lua global shutdown function exists
+  ts_lua_main_ctx *shutdown_main_ctx = &ts_lua_g_main_ctx_array[0];
+  ts_lua_http_ctx *shutdown_http_ctx = ts_lua_create_http_ctx(shutdown_main_ctx, conf);
+  lua_State       *sl                = shutdown_http_ctx->cinfo.routine.lua;
+
+  lua_getglobal(sl, TS_LUA_FUNCTION_G_SHUT_DOWN);
+  if (lua_type(sl, -1) == LUA_TFUNCTION) {
+    TSMutex shutdown_mutex = TSMutexCreate();
+    TSCont  shutdown_contp = TSContCreate(shutdownHookHandler, shutdown_mutex);
+    if (!shutdown_contp) {
+      TSError("[ts_lua][%s] could not create shutdown continuation", __FUNCTION__);
+      if (shutdown_mutex) {
+        TSMutexDestroy(shutdown_mutex);
+      }
+    } else {
+      TSContDataSet(shutdown_contp, conf);
+      TSLifecycleHookAdd(TS_LIFECYCLE_SHUTDOWN_HOOK, shutdown_contp);
+      Dbg(dbg_ctl, "shutdown_hook added");
+    }
+  }
+  lua_pop(sl, 1);
+
+  ts_lua_destroy_http_ctx(shutdown_http_ctx);
+
   // support for reload as global plugin
   if (reload) {
     TSCont config_contp = TSContCreate(configHandler, nullptr);

plugins/lua/ts_lua_common.h

@@ -62,6 +62,9 @@ extern "C" {
 // TLS hooks can only be global
 #define TS_LUA_FUNCTION_G_VCONN_START "do_global_vconn_start"
 
+// Lifecycle hooks
+#define TS_LUA_FUNCTION_G_SHUT_DOWN "__shutdown__"
+
 #define TS_LUA_DEBUG_TAG "ts_lua"
 
 #define TS_LUA_EVENT_COROUTINE_CONT 20000

tests/gold_tests/pluginTest/lua/global_shutdown.lua

@@ -0,0 +1,23 @@
+--  Licensed to the Apache Software Foundation (ASF) under one
+--  or more contributor license agreements.  See the NOTICE file
+--  distributed with this work for additional information
+--  regarding copyright ownership.  The ASF licenses this file
+--  to you under the Apache License, Version 2.0 (the
+--  "License"); you may not use this file except in compliance
+--  with the License.  You may obtain a copy of the License at
+--
+--  http://www.apache.org/licenses/LICENSE-2.0
+--
+--  Unless required by applicable law or agreed to in writing, software
+--  distributed under the License is distributed on an "AS IS" BASIS,
+--  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+--  See the License for the specific language governing permissions and
+--  limitations under the License.
+
+function do_global_read_request()
+  ts.debug('do_global_read_request called')
+end
+
+function __shutdown__()
+  ts.debug('__shutdown__ called')
+end

tests/gold_tests/pluginTest/lua/lua_global_shutdown.test.py

@@ -0,0 +1,64 @@
+'''
+Test __shutdown__ lua global plugin hook.
+'''
+#  Licensed to the Apache Software Foundation (ASF) under one
+#  or more contributor license agreements.  See the NOTICE file
+#  distributed with this work for additional information
+#  regarding copyright ownership.  The ASF licenses this file
+#  to you under the Apache License, Version 2.0 (the
+#  "License"); you may not use this file except in compliance
+#  with the License.  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+Test.Summary = '''
+Test __shutdown__ lua global plugin hook
+'''
+
+Test.SkipUnless(Condition.PluginExists('tslua.so'),)
+
+Test.ContinueOnFail = True
+
+server = Test.MakeOriginServer("server")
+ts = Test.MakeATSProcess("ts")
+
+Test.Setup.Copy("global_shutdown.lua")
+
+request_header = {"headers": "GET / HTTP/1.1\r\nHost: www.example.com\r\n\r\n", "timestamp": "1469733493.993", "body": ""}
+response_header = {"headers": "HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n", "timestamp": "1469733493.993", "body": ""}
+server.addResponse("sessionfile.log", request_header, response_header)
+
+ts.Disk.remap_config.AddLine('map / http://127.0.0.1:{}/'.format(server.Variables.Port))
+
+# Use 2 states so the shutdown handler is called a predictable number of times.
+ts.Disk.plugin_config.AddLine('tslua.so --states=2 {}/global_shutdown.lua'.format(Test.RunDirectory))
+
+ts.Disk.records_config.update({
+    'proxy.config.diags.debug.enabled': 1,
+    'proxy.config.diags.debug.tags': 'ts_lua',
+})
+
+curl_and_args = '-s -D /dev/stdout -o /dev/stderr -x localhost:{} '.format(ts.Variables.port)
+
+# 0 Test - Send a request to confirm the global plugin is active.
+tr = Test.AddTestRun("Lua global read request hook fires for HTTP requests")
+ps = tr.Processes.Default
+ps.StartBefore(server, ready=When.PortOpen(server.Variables.Port))
+ps.StartBefore(Test.Processes.ts)
+tr.MakeCurlCommand(curl_and_args + 'http://www.example.com/', ts=ts)
+ps.ReturnCode = 0
+tr.StillRunningAfter = ts
+
+# Verify do_global_read_request was invoked for the HTTP request above.
+ts.Disk.traffic_out.Content = Testers.ContainsExpression(
+    r'do_global_read_request called', 'do_global_read_request should be called for HTTP requests')
+
+# After all test runs complete AuTest stops ATS, which fires TS_LIFECYCLE_SHUTDOWN_HOOK.
+# The shutdown handler calls __shutdown__ once per Lua state (2 states configured).
+ts.Disk.traffic_out.Content += Testers.ContainsExpression(r'__shutdown__ called', '__shutdown__ should be called on ATS shutdown')