Send end trace on early shutdown

inc/namespace.php

@@ -25,11 +25,16 @@ function bootstrap() {
 	}
 
 	if ( function_exists( 'add_action' ) ) {
-		add_action( 'shutdown', __NAMESPACE__ . '\\on_shutdown' );
-	} else {
-		register_shutdown_function( __NAMESPACE__ . '\\on_shutdown' );
+		// Hook X-Ray into the 'shutdown' action if possible. This allows plugins to control the
+		// load order of shutdown functions as they can use the action's priority argument.
+		add_action( 'shutdown', __NAMESPACE__ . '\\on_shutdown_action' );
 	}
 
+	// As well as using the shutdown action, we register as "direct" shutdown function in the event
+	// that the 'shutdown' action is never registered (such as requests that hit advanced-cache.php),
+	// of is the add_action API is not yet available.
+	register_shutdown_function( __NAMESPACE__ . '\\on_shutdown' );
+
 	$current_error_handler = set_error_handler( function () use ( &$current_error_handler ) { // @codingStandardsIgnoreLine
 		call_user_func_array( __NAMESPACE__ . '\\error_handler', func_get_args() );
 		if ( $current_error_handler ) {
@@ -42,8 +47,10 @@ function bootstrap() {
 
 /**
  * Shutdown callback to process the trace once everything has finished.
+ *
+ * This is called by the 'shutdown' WordPress action.
  */
-function on_shutdown() {
+function on_shutdown_action() {
 	$use_fastcgi_finish_request = function_exists( 'fastcgi_finish_request' );
 	if ( function_exists( 'apply_filters' ) ) {
 		$use_fastcgi_finish_request = apply_filters( 'aws_xray.use_fastcgi_finish_request', $use_fastcgi_finish_request );
@@ -66,6 +73,35 @@ function on_shutdown() {
 	send_trace_to_daemon( get_end_trace() );
 }
 
+/**
+ * Shutdown callback that is triggered via register_shutdown_function.
+ *
+ * In some cases the on_shutdown_action() function will not be called, so we have
+ * this failsafe shutdown function to catch any cases where on_shutdown_action() is not called.
+ *
+ */
+function on_shutdown() {
+	// If we shutdown before the plugin API has loaded, return early.
+	if ( ! function_exists( 'has_action' ) ) {
+		on_shutdown_action();
+		return;
+	}
+	// If the "shutdown_action_hook" function has not been registered
+	// to the shutdown action, call it now.
+	if ( has_action( 'shutdown', __NAMESPACE__ . '\\on_shutdown_action' ) === false ) {
+		on_shutdown_action();
+		return;
+	}
+
+	// It's possible the script is shutting down before register_shutdown_function( 'shutdown_action_hook' )
+	// has been called by WordPress. There's no way to check if this callback has been registered, so we use a heuristic that
+	// get_locale() has been defined, which happens just after WordPress calls register_shutdown_function.
+	if ( ! function_exists( 'get_locale' ) ) {
+		on_shutdown_action();
+		return;
+	}
+}
+
 function error_handler( int $errno, string $errstr, string $errfile = null, int $errline = null ) : bool {
 	global $hm_platform_xray_errors;
 

inc/query_monitor/namespace.php

@@ -28,8 +28,8 @@ function register_qm_collector( array $collectors ) : array {
 	add_filter( 'aws_xray.use_fastcgi_finish_request', '__return_false' );
 
 	// Make sure the XRay shutdown function runs before the Query Monitor one.
-	remove_filter( 'shutdown', 'HM\\Platform\\XRay\\on_shutdown' );
-	add_filter( 'shutdown', 'HM\\Platform\\XRay\\on_shutdown', -1 );
+	remove_filter( 'shutdown', 'HM\\Platform\\XRay\\on_shutdown_action' );
+	add_filter( 'shutdown', 'HM\\Platform\\XRay\\on_shutdown_action', -1 );
 
 	require_once __DIR__ . '/class-collector.php';
 	$collectors['aws-xray'] = new Collector();