From b991373784ed46dba7edb479f0f41461bef21d11 Mon Sep 17 00:00:00 2001
From: Daniel Bachhuber <daniel@handbuilt.co>
Date: Fri, 23 Sep 2016 06:33:30 -0700
Subject: [PATCH] Fix profiling of early hooks

When we throw an Exception to bail out of the call, this means the rest
of the bootstrap code hasn't executed, and important things aren't set
up. Instead, we need to summarize early.
---
 features/profile-hook.feature |  8 ++++++++
 inc/class-command.php         | 23 +++++++++++++++++++----
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/features/profile-hook.feature b/features/profile-hook.feature
index ba8432a..289c98c 100644
--- a/features/profile-hook.feature
+++ b/features/profile-hook.feature
@@ -1,5 +1,13 @@
 Feature: Profile a specific hook
 
+  Scenario: Profile a hook before the template is loaded
+    Given a WP install
+
+    When I run `wp profile --hook=plugins_loaded --fields=callback`
+    Then STDOUT should be a table containing rows:
+      | callback          |
+    And STDERR should be empty
+
   Scenario: Profile a hook without any callbacks
     Given a WP install
 
diff --git a/inc/class-command.php b/inc/class-command.php
index 3913cc7..90fc95f 100644
--- a/inc/class-command.php
+++ b/inc/class-command.php
@@ -14,6 +14,8 @@ class Command {
 	private $current_filter_callbacks = array();
 	private $focus_query_offset = 0;
 
+	private static $exception_message = "Need to bail, because can't restore the hooks";
+
 	/**
 	 * Quickly identify what's slow with WordPress.
 	 *
@@ -141,7 +143,7 @@ class Command {
 			$this->current_filter_callbacks = $wp_filter[ $current_filter ];
 			unset( $wp_filter[ $current_filter ] );
 			call_user_func_array( array( $this, 'do_action' ), func_get_args() );
-			throw new \Exception( "Need to bail, because can't restore the hooks" );
+			throw new \Exception( self::$exception_message );
 		}
 
 		WP_CLI::add_wp_hook( $current_filter, array( $this, 'wp_hook_end' ), 999 );
@@ -269,7 +271,11 @@ class Command {
 		try {
 			WP_CLI::get_runner()->load_wordpress();
 		} catch( \Exception $e ) {
-			// pass through
+			// If this was thrown by our do_action implementation, then we need to bail
+			if ( self::$exception_message === $e->getMessage() ) {
+				return;
+			}
+			// Otherwise, pass through.
 		}
 		if ( isset( $this->loggers['wp_profile_last_hook'] ) && $this->loggers['wp_profile_last_hook']->running() ) {
 			$this->loggers['wp_profile_last_hook']->stop();
@@ -295,7 +301,11 @@ class Command {
 		try {
 			wp();
 		} catch( \Exception $e ) {
-			// pass through
+			// If this was thrown by our do_action implementation, then we need to bail
+			if ( self::$exception_message === $e->getMessage() ) {
+				return;
+			}
+			// Otherwise, pass through.
 		}
 		if ( isset( $this->loggers['wp_profile_last_hook'] ) && $this->loggers['wp_profile_last_hook']->running() ) {
 			$this->loggers['wp_profile_last_hook']->stop();
@@ -330,7 +340,12 @@ class Command {
 		try {
 			require_once( ABSPATH . WPINC . '/template-loader.php' );
 		} catch( \Exception $e ) {
-			// pass through
+			// If this was thrown by our do_action implementation, then we need to bail
+			if ( self::$exception_message === $e->getMessage() ) {
+				ob_get_clean();
+				return;
+			}
+			// Otherwise, pass through.
 		}
 		if ( isset( $this->loggers['wp_profile_last_hook'] ) && $this->loggers['wp_profile_last_hook']->running() ) {
 			$this->loggers['wp_profile_last_hook']->stop();