From be748915a63468487a2d28a0e6457280d8f8805b Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Wed, 16 Dec 2015 15:46:43 -0800
Subject: [PATCH] Set TransactionProfiler expectations in the API

This is similar to the main() method for index.php but also
makes use of write mode checks to more granular.

Change-Id: I40fbaef0f5eaefaf2c6d0937d49016a17da532da
---
 includes/DefaultSettings.php |  5 +++++
 includes/api/ApiMain.php     | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/includes/DefaultSettings.php b/includes/DefaultSettings.php
index 0d4ecacb0e..a6d34b0d35 100644
--- a/includes/DefaultSettings.php
+++ b/includes/DefaultSettings.php
@@ -5556,6 +5556,11 @@ $wgTrxProfilerLimits = array(
 		'writeQueryTime' => 1,
 		'maxAffected' => 500
 	),
+	'POST-nonwrite' => array(
+		'masterConns' => 0,
+		'writes' => 0,
+		'readQueryTime' => 5
+	),
 	// Background job runner
 	'JobRunner' => array(
 		'readQueryTime' => 30,
diff --git a/includes/api/ApiMain.php b/includes/api/ApiMain.php
index 49b9786ccb..89ff19a196 100644
--- a/includes/api/ApiMain.php
+++ b/includes/api/ApiMain.php
@@ -1253,6 +1253,8 @@ class ApiMain extends ApiBase {
 		$module = $this->setupModule();
 		$this->mModule = $module;
 
+		$this->setRequestExpectations( $module );
+
 		$this->checkExecutePermissions( $module );
 
 		if ( !$this->checkMaxLag( $module, $params ) ) {
@@ -1284,6 +1286,24 @@ class ApiMain extends ApiBase {
 		}
 	}
 
+	/**
+	 * Set database connection, query, and write expectations given this module request
+	 * @param ApiBase $module
+	 */
+	protected function setRequestExpectations( ApiBase $module ) {
+		$limits = $this->getConfig()->get( 'TrxProfilerLimits' );
+		$trxProfiler = Profiler::instance()->getTransactionProfiler();
+		if ( $this->getRequest()->wasPosted() ) {
+			if ( $module->isWriteMode() ) {
+				$trxProfiler->setExpectations( $limits['POST'], __METHOD__ );
+			} else {
+				$trxProfiler->setExpectations( $limits['POST-nonwrite'], __METHOD__ );
+			}
+		} else {
+			$trxProfiler->setExpectations( $limits['GET'], __METHOD__ );
+		}
+	}
+
 	/**
 	 * Log the preceding request
 	 * @param float $time Time in seconds
-- 
2.20.1