From e2e2eb49d6581e2eab58ba843ef2db84ce06b256 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Sat, 29 Jun 2019 01:36:01 -0700
Subject: [PATCH] objectcache: optimize lock() and unlock() methods in
 SqlBagOStuff

Also clean up base method in BagOStuff.

Bug: T113916
Change-Id: I3a1c6afe531a8eae34608bc7fe0aa6f9f4d439fe
---
 includes/libs/objectcache/BagOStuff.php    |  6 +-
 includes/libs/rdbms/database/IDatabase.php |  2 +-
 includes/objectcache/SqlBagOStuff.php      | 73 +++++++++++++++++++++-
 3 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/includes/libs/objectcache/BagOStuff.php b/includes/libs/objectcache/BagOStuff.php
index 50441c5e31..7759947ee0 100644
--- a/includes/libs/objectcache/BagOStuff.php
+++ b/includes/libs/objectcache/BagOStuff.php
@@ -590,7 +590,11 @@ abstract class BagOStuff implements IExpiringStore, LoggerAwareInterface {
 	 * @return bool Success
 	 */
 	public function unlock( $key ) {
-		if ( isset( $this->locks[$key] ) && --$this->locks[$key]['depth'] <= 0 ) {
+		if ( !isset( $this->locks[$key] ) ) {
+			return false;
+		}
+
+		if ( --$this->locks[$key]['depth'] <= 0 ) {
 			unset( $this->locks[$key] );
 
 			$ok = $this->doDelete( "{$key}:lock" );
diff --git a/includes/libs/rdbms/database/IDatabase.php b/includes/libs/rdbms/database/IDatabase.php
index fca2c005e3..ad6d4d2f45 100644
--- a/includes/libs/rdbms/database/IDatabase.php
+++ b/includes/libs/rdbms/database/IDatabase.php
@@ -2112,7 +2112,7 @@ interface IDatabase {
 	 *
 	 * @param string $lockName Name of lock to aquire
 	 * @param string $method Name of the calling method
-	 * @param int $timeout Acquisition timeout in seconds
+	 * @param int $timeout Acquisition timeout in seconds (0 means non-blocking)
 	 * @return bool
 	 * @throws DBError
 	 */
diff --git a/includes/objectcache/SqlBagOStuff.php b/includes/objectcache/SqlBagOStuff.php
index 8a5f59163a..39d0353d45 100644
--- a/includes/objectcache/SqlBagOStuff.php
+++ b/includes/objectcache/SqlBagOStuff.php
@@ -338,6 +338,7 @@ class SqlBagOStuff extends BagOStuff {
 
 		$result = true;
 		$exptime = (int)$expiry;
+		/** @noinspection PhpUnusedLocalVariableInspection */
 		$silenceScope = $this->silenceTransactionProfiler();
 		foreach ( $keysByTable as $serverIndex => $serverKeys ) {
 			$db = null;
@@ -406,6 +407,7 @@ class SqlBagOStuff extends BagOStuff {
 	protected function cas( $casToken, $key, $value, $exptime = 0, $flags = 0 ) {
 		list( $serverIndex, $tableName ) = $this->getTableByKey( $key );
 		$db = null;
+		/** @noinspection PhpUnusedLocalVariableInspection */
 		$silenceScope = $this->silenceTransactionProfiler();
 		try {
 			$db = $this->getDB( $serverIndex );
@@ -457,6 +459,7 @@ class SqlBagOStuff extends BagOStuff {
 		}
 
 		$result = true;
+		/** @noinspection PhpUnusedLocalVariableInspection */
 		$silenceScope = $this->silenceTransactionProfiler();
 		foreach ( $keysByTable as $serverIndex => $serverKeys ) {
 			$db = null;
@@ -495,6 +498,7 @@ class SqlBagOStuff extends BagOStuff {
 	public function incr( $key, $step = 1 ) {
 		list( $serverIndex, $tableName ) = $this->getTableByKey( $key );
 		$db = null;
+		/** @noinspection PhpUnusedLocalVariableInspection */
 		$silenceScope = $this->silenceTransactionProfiler();
 		try {
 			$db = $this->getDB( $serverIndex );
@@ -553,6 +557,7 @@ class SqlBagOStuff extends BagOStuff {
 	public function changeTTL( $key, $exptime = 0, $flags = 0 ) {
 		list( $serverIndex, $tableName ) = $this->getTableByKey( $key );
 		$db = null;
+		/** @noinspection PhpUnusedLocalVariableInspection */
 		$silenceScope = $this->silenceTransactionProfiler();
 		try {
 			$db = $this->getDB( $serverIndex );
@@ -635,6 +640,7 @@ class SqlBagOStuff extends BagOStuff {
 	 * @return bool
 	 */
 	public function deleteObjectsExpiringBefore( $timestamp, $progressCallback = false ) {
+		/** @noinspection PhpUnusedLocalVariableInspection */
 		$silenceScope = $this->silenceTransactionProfiler();
 		for ( $serverIndex = 0; $serverIndex < $this->numServers; $serverIndex++ ) {
 			$db = null;
@@ -713,6 +719,7 @@ class SqlBagOStuff extends BagOStuff {
 	 * @return bool
 	 */
 	public function deleteAll() {
+		/** @noinspection PhpUnusedLocalVariableInspection */
 		$silenceScope = $this->silenceTransactionProfiler();
 		for ( $serverIndex = 0; $serverIndex < $this->numServers; $serverIndex++ ) {
 			$db = null;
@@ -729,6 +736,68 @@ class SqlBagOStuff extends BagOStuff {
 		return true;
 	}
 
+	public function lock( $key, $timeout = 6, $expiry = 6, $rclass = '' ) {
+		// Avoid deadlocks and allow lock reentry if specified
+		if ( isset( $this->locks[$key] ) ) {
+			if ( $rclass != '' && $this->locks[$key]['class'] === $rclass ) {
+				++$this->locks[$key]['depth'];
+				return true;
+			} else {
+				return false;
+			}
+		}
+
+		list( $serverIndex ) = $this->getTableByKey( $key );
+		try {
+			$db = $this->getDB( $serverIndex );
+			$ok = $db->lock( $key, __METHOD__, $timeout );
+			if ( $ok ) {
+				$this->locks[$key] = [ 'class' => $rclass, 'depth' => 1 ];
+			}
+
+			$this->logger->warning(
+				__METHOD__ . " failed due to timeout for {key}.",
+				[ 'key' => $key, 'timeout' => $timeout ]
+			);
+
+			return $ok;
+		} catch ( DBError $e ) {
+			$this->handleWriteError( $e, $db, $serverIndex );
+			$ok = false;
+		}
+
+		return $ok;
+	}
+
+	public function unlock( $key ) {
+		if ( !isset( $this->locks[$key] ) ) {
+			return false;
+		}
+
+		if ( --$this->locks[$key]['depth'] <= 0 ) {
+			unset( $this->locks[$key] );
+
+			list( $serverIndex ) = $this->getTableByKey( $key );
+			try {
+				$db = $this->getDB( $serverIndex );
+				$ok = $db->unlock( $key, __METHOD__ );
+				if ( !$ok ) {
+					$this->logger->warning(
+						__METHOD__ . ' failed to release lock for {key}.',
+						[ 'key' => $key ]
+					);
+				}
+			} catch ( DBError $e ) {
+				$this->handleWriteError( $e, $db, $serverIndex );
+				$ok = false;
+			}
+		} else {
+			$ok = false;
+		}
+
+		return $ok;
+	}
+
 	/**
 	 * Serialize an object and, if possible, compress the representation.
 	 * On typical message and page data, this can provide a 3X decrease
@@ -790,8 +859,8 @@ class SqlBagOStuff extends BagOStuff {
 	 * @param int $serverIndex
 	 * @throws Exception
 	 */
-	protected function handleWriteError( DBError $exception, IDatabase $db = null, $serverIndex ) {
-		if ( !$db ) {
+	protected function handleWriteError( DBError $exception, $db, $serverIndex ) {
+		if ( !( $db instanceof IDatabase ) ) {
 			$this->markServerDown( $exception, $serverIndex );
 		}
 
-- 
2.20.1