From b31f5c39faa71447632fc9233471eff691b644d0 Mon Sep 17 00:00:00 2001
From: Daniel Bachhuber <daniel@handbuilt.co>
Date: Thu, 8 Sep 2016 17:42:58 -0700
Subject: [PATCH] Ensure the last step to a stage is logged appropriately

---
 features/profile-stage-template.feature |  1 +
 inc/class-command.php                   | 13 +++++++++++++
 inc/class-logger.php                    |  7 +++++++
 3 files changed, 21 insertions(+)

diff --git a/features/profile-stage-template.feature b/features/profile-stage-template.feature
index f255775..35b065c 100644
--- a/features/profile-stage-template.feature
+++ b/features/profile-stage-template.feature
@@ -18,4 +18,5 @@ Feature: Profile the template render stage
       | loop_end          |
       |                   |
       | wp_footer         |
+      |                   |
       | total             |
diff --git a/inc/class-command.php b/inc/class-command.php
index 6e1ce27..3913cc7 100644
--- a/inc/class-command.php
+++ b/inc/class-command.php
@@ -217,6 +217,10 @@ class Command {
 				$pseudo_hook = "before {$this->stage_hooks[$key+1]}";
 				$this->loggers[ $pseudo_hook ] = new Logger( 'hook', '' );
 				$this->loggers[ $pseudo_hook ]->start();
+			} else {
+				$pseudo_hook = 'wp_profile_last_hook';
+				$this->loggers[ $pseudo_hook ] = new Logger( 'hook', '' );
+				$this->loggers[ $pseudo_hook ]->start();
 			}
 		}
 
@@ -267,6 +271,9 @@ class Command {
 		} catch( \Exception $e ) {
 			// pass through
 		}
+		if ( isset( $this->loggers['wp_profile_last_hook'] ) && $this->loggers['wp_profile_last_hook']->running() ) {
+			$this->loggers['wp_profile_last_hook']->stop();
+		}
 		if ( ! $this->focus_stage && ! $this->focus_hook ) {
 			$logger->stop();
 			$this->loggers[] = $logger;
@@ -290,6 +297,9 @@ class Command {
 		} catch( \Exception $e ) {
 			// pass through
 		}
+		if ( isset( $this->loggers['wp_profile_last_hook'] ) && $this->loggers['wp_profile_last_hook']->running() ) {
+			$this->loggers['wp_profile_last_hook']->stop();
+		}
 		if ( ! $this->focus_stage && ! $this->focus_hook ) {
 			$logger->stop();
 			$this->loggers[] = $logger;
@@ -322,6 +332,9 @@ class Command {
 		} catch( \Exception $e ) {
 			// pass through
 		}
+		if ( isset( $this->loggers['wp_profile_last_hook'] ) && $this->loggers['wp_profile_last_hook']->running() ) {
+			$this->loggers['wp_profile_last_hook']->stop();
+		}
 		ob_get_clean();
 		if ( ! $this->focus_stage && ! $this->focus_hook ) {
 			$logger->stop();
diff --git a/inc/class-logger.php b/inc/class-logger.php
index 8213110..77a08bc 100644
--- a/inc/class-logger.php
+++ b/inc/class-logger.php
@@ -43,6 +43,13 @@ class Logger {
 		$this->cache_miss_offset = ! empty( $wp_object_cache->cache_misses ) ? $wp_object_cache->cache_misses : 0;
 	}
 
+	/**
+	 * Whether or not the logger is running
+	 */
+	public function running() {
+		return ! is_null( $this->start_time );
+	}
+
 	/**
 	 * Stop this logger
 	 */