Add a lock to storage upgrade and adjustment

Summary: Fixes T9715. Adds a MySQL-based lock to ensure that schema migrations are not applied on multiple hosts simultaneously.

Test Plan: Ran `./bin/storage upgrade` concurrently. One invocation was successful whilst the other hit a `PhutilLockException`.

Reviewers: #blessed_reviewers, epriestley

Subscribers: Korvin

Maniphest Tasks: T9715

Differential Revision: https://secure.phabricator.com/D14463

scripts/sql/manage_storage.php

@@ -70,9 +70,10 @@ try {
       array(
         'name'    => 'disable-utf8mb4',
         'help'    => pht(
-          'Disable utf8mb4, even if the database supports it. This is an '.
+          'Disable %s, even if the database supports it. This is an '.
           'advanced feature used for testing changes to Phabricator; you '.
-          'should not normally use this flag.'),
+          'should not normally use this flag.',
+          'utf8mb4'),
       ),
     ));
 } catch (PhutilArgumentUsageException $ex) {
@@ -83,12 +84,12 @@ try {
 // First, test that the Phabricator configuration is set up correctly. After
 // we know this works we'll test any administrative credentials specifically.
 
-$test_api = new PhabricatorStorageManagementAPI();
+$test_api = id(new PhabricatorStorageManagementAPI())
-$test_api->setUser($default_user);
+  ->setUser($default_user)
-$test_api->setHost($default_host);
+  ->setHost($default_host)
-$test_api->setPort($default_port);
+  ->setPort($default_port)
-$test_api->setPassword($conf->getPassword());
+  ->setPassword($conf->getPassword())
-$test_api->setNamespace($args->getArg('namespace'));
+  ->setNamespace($args->getArg('namespace'));
 
 try {
   queryfx(
@@ -113,13 +114,10 @@ try {
       '--password'),
     pht('Raw MySQL Error'),
     $ex->getMessage());
-
   echo phutil_console_wrap($message);
-
   exit(1);
 }
 
-
 if ($args->getArg('password') === null) {
   // This is already a PhutilOpaqueEnvelope.
   $password = $conf->getPassword();
@@ -129,14 +127,14 @@ if ($args->getArg('password') === null) {
   PhabricatorEnv::overrideConfig('mysql.pass', $args->getArg('password'));
 }
 
-$api = new PhabricatorStorageManagementAPI();
+$api = id(new PhabricatorStorageManagementAPI())
-$api->setUser($args->getArg('user'));
+  ->setUser($args->getArg('user'))
-PhabricatorEnv::overrideConfig('mysql.user', $args->getArg('user'));
+  ->setHost($default_host)
-$api->setHost($default_host);
+  ->setPort($default_port)
-$api->setPort($default_port);
+  ->setPassword($password)
-$api->setPassword($password);
+  ->setNamespace($args->getArg('namespace'))
-$api->setNamespace($args->getArg('namespace'));
+  ->setDisableUTF8MB4($args->getArg('disable-utf8mb4'));
-$api->setDisableUTF8MB4($args->getArg('disable-utf8mb4'));
+PhabricatorEnv::overrideConfig('mysql.user', $api->getUser());
 
 try {
   queryfx(
@@ -154,9 +152,7 @@ try {
       '--password'),
     pht('Raw MySQL Error'),
     $ex->getMessage());
-
   echo phutil_console_wrap($message);
-
   exit(1);
 }
 

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementAdjustWorkflow.php

@@ -24,13 +24,11 @@ final class PhabricatorStorageManagementAdjustWorkflow
         ));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
-    $force = $args->getArg('force');
     $unsafe = $args->getArg('unsafe');
-    $dry_run = $args->getArg('dryrun');
 
     $this->requireAllPatchesApplied();
-    return $this->adjustSchemata($force, $unsafe, $dry_run);
+    return $this->adjustSchemata($unsafe);
   }
 
   private function requireAllPatchesApplied() {

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementDatabasesWorkflow.php

@@ -10,13 +10,12 @@ final class PhabricatorStorageManagementDatabasesWorkflow
       ->setSynopsis(pht('List Phabricator databases.'));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
     $api     = $this->getAPI();
     $patches = $this->getPatches();
 
-    $databases = $api->getDatabaseList($patches, $only_living = true);
+    $databases = $api->getDatabaseList($patches, true);
     echo implode("\n", $databases)."\n";
-
     return 0;
   }
 

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementDestroyWorkflow.php

@@ -20,24 +20,24 @@ final class PhabricatorStorageManagementDestroyWorkflow
         ));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
-    $is_dry = $args->getArg('dryrun');
+    $console = PhutilConsole::getConsole();
-    $is_force = $args->getArg('force');
 
-    if (!$is_dry && !$is_force) {
+    if (!$this->isDryRun() && !$this->isForce()) {
-      echo phutil_console_wrap(
+      $console->writeOut(
+        phutil_console_wrap(
           pht(
-          'Are you completely sure you really want to permanently destroy all '.
+            'Are you completely sure you really want to permanently destroy '.
-          'storage for Phabricator data? This operation can not be undone and '.
+            'all storage for Phabricator data? This operation can not be '.
-          'your data will not be recoverable if you proceed.'));
+            'undone and your data will not be recoverable if you proceed.')));
 
       if (!phutil_console_confirm(pht('Permanently destroy all data?'))) {
-        echo pht('Cancelled.')."\n";
+        $console->writeOut("%s\n", pht('Cancelled.'));
         exit(1);
       }
 
       if (!phutil_console_confirm(pht('Really destroy all data forever?'))) {
-        echo pht('Cancelled.')."\n";
+        $console->writeOut("%s\n", pht('Cancelled.'));
         exit(1);
       }
     }
@@ -57,16 +57,21 @@ final class PhabricatorStorageManagementDestroyWorkflow
     } else {
       $databases   = $api->getDatabaseList($patches);
       $databases[] = $api->getDatabaseName('meta_data');
+
       // These are legacy databases that were dropped long ago. See T2237.
       $databases[] = $api->getDatabaseName('phid');
       $databases[] = $api->getDatabaseName('directory');
     }
 
     foreach ($databases as $database) {
-      if ($is_dry) {
+      if ($this->isDryRun()) {
-        echo pht("DRYRUN: Would drop database '%s'.", $database)."\n";
+        $console->writeOut(
+          "%s\n",
+          pht("DRYRUN: Would drop database '%s'.", $database));
       } else {
-        echo pht("Dropping database '%s'...", $database)."\n";
+        $console->writeOut(
+          "%s\n",
+          pht("Dropping database '%s'...", $database));
         queryfx(
           $api->getConn(null),
           'DROP DATABASE IF EXISTS %T',
@@ -74,8 +79,8 @@ final class PhabricatorStorageManagementDestroyWorkflow
       }
     }
 
-    if (!$is_dry) {
+    if (!$this->isDryRun()) {
-      echo pht('Storage was destroyed.')."\n";
+      $console->writeOut("%s\n", pht('Storage was destroyed.'));
     }
 
     return 0;

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementDumpWorkflow.php

@@ -10,11 +10,12 @@ final class PhabricatorStorageManagementDumpWorkflow
       ->setSynopsis(pht('Dump all data in storage to stdout.'));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
-    $console = PhutilConsole::getConsole();
     $api     = $this->getAPI();
     $patches = $this->getPatches();
 
+    $console = PhutilConsole::getConsole();
+
     $applied = $api->getAppliedPatches();
     if ($applied === null) {
       $namespace = $api->getNamespace();
@@ -24,11 +25,11 @@ final class PhabricatorStorageManagementDumpWorkflow
           'initialized in this storage namespace ("%s"). Use '.
           '**%s** to initialize storage.',
           $namespace,
-          'storage upgrade'));
+          './bin/storage upgrade'));
       return 1;
     }
 
-    $databases = $api->getDatabaseList($patches, $only_living = true);
+    $databases = $api->getDatabaseList($patches, true);
 
     list($host, $port) = $this->getBareHostAndPort($api->getHost());
 

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementProbeWorkflow.php

@@ -10,7 +10,7 @@ final class PhabricatorStorageManagementProbeWorkflow
       ->setSynopsis(pht('Show approximate table sizes.'));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
     $console = PhutilConsole::getConsole();
     $console->writeErr(
       "%s\n",
@@ -18,7 +18,7 @@ final class PhabricatorStorageManagementProbeWorkflow
 
     $api       = $this->getAPI();
     $patches   = $this->getPatches();
-    $databases = $api->getDatabaseList($patches, $only_living = true);
+    $databases = $api->getDatabaseList($patches, true);
 
     $conn_r = $api->getConn(null);
 

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementQuickstartWorkflow.php

@@ -22,6 +22,8 @@ final class PhabricatorStorageManagementQuickstartWorkflow
   }
 
   public function execute(PhutilArgumentParser $args) {
+    parent::execute($args);
+
     $output = $args->getArg('output');
     if (!$output) {
       throw new PhutilArgumentUsageException(
@@ -38,8 +40,10 @@ final class PhabricatorStorageManagementQuickstartWorkflow
       throw new PhutilArgumentUsageException(
         pht(
           'You can only generate a new quickstart file if MySQL supports '.
-          'the utf8mb4 character set (available in MySQL 5.5 and newer). The '.
+          'the %s character set (available in MySQL 5.5 and newer). The '.
-          'configured server does not support utf8mb4.'));
+          'configured server does not support %s.',
+          'utf8mb4',
+          'utf8mb4'));
     }
 
     $err = phutil_passthru(
@@ -139,7 +143,7 @@ final class PhabricatorStorageManagementQuickstartWorkflow
     $dump = preg_replace('/^--.*$/m', '', $dump);
 
     // Remove table drops, locks, and unlocks. These are never relevant when
-    // performing q quickstart.
+    // performing a quickstart.
     $dump = preg_replace(
       '/^(DROP TABLE|LOCK TABLES|UNLOCK TABLES).*$/m',
       '',

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementRenamespaceWorkflow.php

@@ -30,25 +30,31 @@ final class PhabricatorStorageManagementRenamespaceWorkflow
         ));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
     $console = PhutilConsole::getConsole();
 
     $in = $args->getArg('in');
     if (!strlen($in)) {
       throw new PhutilArgumentUsageException(
-        pht('Specify the dumpfile to read with --in.'));
+        pht(
+          'Specify the dumpfile to read with %s.',
+          '--in'));
     }
 
     $from = $args->getArg('from');
     if (!strlen($from)) {
       throw new PhutilArgumentUsageException(
-        pht('Specify namespace to rename from with --from.'));
+        pht(
+          'Specify namespace to rename from with %s.',
+          '--from'));
     }
 
     $to = $args->getArg('to');
     if (!strlen($to)) {
       throw new PhutilArgumentUsageException(
-        pht('Specify namespace to rename to with --to.'));
+        pht(
+          'Specify namespace to rename to with %s.',
+          '--to'));
     }
 
     $patterns = array(

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementShellWorkflow.php

@@ -11,6 +11,8 @@ final class PhabricatorStorageManagementShellWorkflow
   }
 
   public function execute(PhutilArgumentParser $args) {
+
+
     $api = $this->getAPI();
     list($host, $port) = $this->getBareHostAndPort($api->getHost());
 

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementStatusWorkflow.php

@@ -10,7 +10,7 @@ final class PhabricatorStorageManagementStatusWorkflow
       ->setSynopsis(pht('Show patch application status.'));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
     $api     = $this->getAPI();
     $patches = $this->getPatches();
 
@@ -20,7 +20,7 @@ final class PhabricatorStorageManagementStatusWorkflow
       echo phutil_console_format(
         "**%s**: %s\n",
         pht('Database Not Initialized'),
-        pht('Run **%s** to initialize.', 'storage upgrade'));
+        pht('Run **%s** to initialize.', './bin/storage upgrade'));
 
       return 1;
     }

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementUpgradeWorkflow.php

@@ -21,7 +21,7 @@ final class PhabricatorStorageManagementUpgradeWorkflow
           array(
             'name'  => 'no-quickstart',
             'help'  => pht(
-              'Build storage patch-by-patch from scatch, even if it could '.
+              'Build storage patch-by-patch from scratch, even if it could '.
               'be loaded from the quickstart template.'),
           ),
           array(
@@ -37,23 +37,21 @@ final class PhabricatorStorageManagementUpgradeWorkflow
         ));
   }
 
-  public function execute(PhutilArgumentParser $args) {
+  public function didExecute(PhutilArgumentParser $args) {
-    $is_dry = $args->getArg('dryrun');
+    $console = PhutilConsole::getConsole();
-    $is_force = $args->getArg('force');
-
-    $api = $this->getAPI();
     $patches = $this->getPatches();
 
-    if (!$is_dry && !$is_force) {
+    if (!$this->isDryRun() && !$this->isForce()) {
-      echo phutil_console_wrap(
+      $console->writeOut(
+        phutil_console_wrap(
           pht(
             'Before running storage upgrades, you should take down the '.
             'Phabricator web interface and stop any running Phabricator '.
             'daemons (you can disable this warning with %s).',
-          '--force'));
+            '--force')));
 
       if (!phutil_console_confirm(pht('Are you ready to continue?'))) {
-        echo pht('Cancelled.')."\n";
+        $console->writeOut("%s\n", pht('Cancelled.'));
         return 1;
       }
     }
@@ -67,7 +65,7 @@ final class PhabricatorStorageManagementUpgradeWorkflow
             "Use '%s' to show patch status.",
             '--apply',
             $apply_only,
-            'storage status'));
+            './bin/storage status'));
       }
     }
 
@@ -75,155 +73,15 @@ final class PhabricatorStorageManagementUpgradeWorkflow
     $init_only     = $args->getArg('init-only');
     $no_adjust     = $args->getArg('no-adjust');
 
-    $applied = $api->getAppliedPatches();
+    $this->upgradeSchemata($apply_only, $no_quickstart, $init_only);
-    if ($applied === null) {
 
-      if ($is_dry) {
-        echo pht(
-          "DRYRUN: Patch metadata storage doesn't exist yet, ".
-          "it would be created.\n");
-        return 0;
-      }
-
-      if ($apply_only) {
-        throw new PhutilArgumentUsageException(
-          pht(
-            'Storage has not been initialized yet, you must initialize '.
-            'storage before selectively applying patches.'));
-        return 1;
-      }
-
-      $legacy = $api->getLegacyPatches($patches);
-      if ($legacy || $no_quickstart || $init_only) {
-
-        // If we have legacy patches, we can't quickstart.
-
-        $api->createDatabase('meta_data');
-        $api->createTable(
-          'meta_data',
-          'patch_status',
-          array(
-            'patch VARCHAR(255) NOT NULL PRIMARY KEY COLLATE utf8_general_ci',
-            'applied INT UNSIGNED NOT NULL',
-          ));
-
-        foreach ($legacy as $patch) {
-          $api->markPatchApplied($patch);
-        }
-      } else {
-        echo pht('Loading quickstart template...')."\n";
-        $root = dirname(phutil_get_library_root('phabricator'));
-        $sql  = $root.'/resources/sql/quickstart.sql';
-        $api->applyPatchSQL($sql);
-      }
-    }
-
-    if ($init_only) {
-      echo pht('Storage initialized.')."\n";
-      return 0;
-    }
-
-    $applied = $api->getAppliedPatches();
-    $applied = array_fuse($applied);
-
-    $skip_mark = false;
-    if ($apply_only) {
-      if (isset($applied[$apply_only])) {
-
-        unset($applied[$apply_only]);
-        $skip_mark = true;
-
-        if (!$is_force && !$is_dry) {
-          echo phutil_console_wrap(
-            pht(
-              "Patch '%s' has already been applied. Are you sure you want ".
-              "to apply it again? This may put your storage in a state ".
-              "that the upgrade scripts can not automatically manage.",
-              $apply_only));
-          if (!phutil_console_confirm(pht('Apply patch again?'))) {
-            echo pht('Cancelled.')."\n";
-            return 1;
-          }
-        }
-      }
-    }
-
-    while (true) {
-      $applied_something = false;
-      foreach ($patches as $key => $patch) {
-        if (isset($applied[$key])) {
-          unset($patches[$key]);
-          continue;
-        }
-
-        if ($apply_only && $apply_only != $key) {
-          unset($patches[$key]);
-          continue;
-        }
-
-        $can_apply = true;
-        foreach ($patch->getAfter() as $after) {
-          if (empty($applied[$after])) {
-            if ($apply_only) {
-              echo pht(
-                "Unable to apply patch '%s' because it depends ".
-                "on patch '%s', which has not been applied.\n",
-                $apply_only,
-                $after);
-              return 1;
-            }
-            $can_apply = false;
-            break;
-          }
-        }
-
-        if (!$can_apply) {
-          continue;
-        }
-
-        $applied_something = true;
-
-        if ($is_dry) {
-          echo pht("DRYRUN: Would apply patch '%s'.", $key)."\n";
-        } else {
-          echo pht("Applying patch '%s'...", $key)."\n";
-
-          $t_begin = microtime(true);
-          $api->applyPatch($patch);
-          $t_end = microtime(true);
-
-          if (!$skip_mark) {
-            $api->markPatchApplied($key, ($t_end - $t_begin));
-          }
-        }
-
-        unset($patches[$key]);
-        $applied[$key] = true;
-      }
-
-      if (!$applied_something) {
-        if (count($patches)) {
-          throw new Exception(
-            pht(
-              'Some patches could not be applied: %s',
-              implode(', ', array_keys($patches))));
-        } else if (!$is_dry && !$apply_only) {
-          echo pht(
-            "Storage is up to date. Use '%s' for details.",
-            'storage status')."\n";
-        }
-        break;
-      }
-    }
-
-    $console = PhutilConsole::getConsole();
     if ($no_adjust || $init_only || $apply_only) {
       $console->writeOut(
         "%s\n",
         pht('Declining to apply storage adjustments.'));
       return 0;
     } else {
-      return $this->adjustSchemata($is_force, $unsafe = false, $is_dry);
+      return $this->adjustSchemata(false);
     }
   }
 

src/infrastructure/storage/management/workflow/PhabricatorStorageManagementWorkflow.php

@@ -3,17 +3,13 @@
 abstract class PhabricatorStorageManagementWorkflow
   extends PhabricatorManagementWorkflow {
 
-  private $patches;
   private $api;
+  private $dryRun;
+  private $force;
+  private $patches;
 
-  public function setPatches(array $patches) {
+  final public function getAPI() {
-    assert_instances_of($patches, 'PhabricatorStoragePatch');
+    return $this->api;
-    $this->patches = $patches;
-    return $this;
-  }
-
-  public function getPatches() {
-    return $this->patches;
   }
 
   final public function setAPI(PhabricatorStorageManagementAPI $api) {
@@ -21,10 +17,44 @@ abstract class PhabricatorStorageManagementWorkflow
     return $this;
   }
 
-  final public function getAPI() {
+  final protected function isDryRun() {
-    return $this->api;
+    return $this->dryRun;
   }
 
+  final protected function setDryRun($dry_run) {
+    $this->dryRun = $dry_run;
+    return $this;
+  }
+
+  final protected function isForce() {
+    return $this->force;
+  }
+
+  final protected function setForce($force) {
+    $this->force = $force;
+    return $this;
+  }
+
+  public function getPatches() {
+    return $this->patches;
+  }
+
+  public function setPatches(array $patches) {
+    assert_instances_of($patches, 'PhabricatorStoragePatch');
+    $this->patches = $patches;
+    return $this;
+  }
+
+
+  public function execute(PhutilArgumentParser $args) {
+    $this->setDryRun($args->getArg('dryrun'));
+    $this->setForce($args->getArg('force'));
+
+    $this->didExecute($args);
+  }
+
+  public function didExecute(PhutilArgumentParser $args) {}
+
   private function loadSchemata() {
     $query = id(new PhabricatorConfigSchemaQuery())
       ->setAPI($this->getAPI());
@@ -36,7 +66,20 @@ abstract class PhabricatorStorageManagementWorkflow
     return array($comp, $expect, $actual);
   }
 
-  protected function adjustSchemata($force, $unsafe, $dry_run) {
+  final protected function adjustSchemata($unsafe) {
+    $lock = $this->lock();
+
+    try {
+      $this->doAdjustSchemata($unsafe);
+    } catch (Exception $ex) {
+      $lock->unlock();
+      throw $ex;
+    }
+
+    $lock->unlock();
+  }
+
+  final private function doAdjustSchemata($unsafe) {
     $console = PhutilConsole::getConsole();
 
     $console->writeOut(
@@ -54,7 +97,7 @@ abstract class PhabricatorStorageManagementWorkflow
       return $this->printErrors($errors, 0);
     }
 
-    if (!$force && !$api->isCharacterSetAvailable('utf8mb4')) {
+    if (!$this->force && !$api->isCharacterSetAvailable('utf8mb4')) {
       $message = pht(
         "You have an old version of MySQL (older than 5.5) which does not ".
         "support the utf8mb4 character set. We strongly recomend upgrading to ".
@@ -110,12 +153,12 @@ abstract class PhabricatorStorageManagementWorkflow
 
     $table->draw();
 
-    if ($dry_run) {
+    if ($this->dryRun) {
       $console->writeOut(
         "%s\n",
         pht('DRYRUN: Would apply adjustments.'));
       return 0;
-    } else if (!$force) {
+    } else if (!$this->force) {
       $console->writeOut(
         "\n%s\n",
         pht(
@@ -665,6 +708,171 @@ abstract class PhabricatorStorageManagementWorkflow
     return 2;
   }
 
+  final protected function upgradeSchemata(
+    $apply_only = null,
+    $no_quickstart = false,
+    $init_only = false) {
+
+    $lock = $this->lock();
+
+    try {
+      $this->doUpgradeSchemata($apply_only, $no_quickstart, $init_only);
+    } catch (Exception $ex) {
+      $lock->unlock();
+      throw $ex;
+    }
+
+    $lock->unlock();
+  }
+
+  final private function doUpgradeSchemata(
+    $apply_only,
+    $no_quickstart,
+    $init_only) {
+
+    $api = $this->getAPI();
+
+    $applied = $this->getApi()->getAppliedPatches();
+    if ($applied === null) {
+      if ($this->dryRun) {
+        echo pht(
+          "DRYRUN: Patch metadata storage doesn't exist yet, ".
+          "it would be created.\n");
+        return 0;
+      }
+
+      if ($apply_only) {
+        throw new PhutilArgumentUsageException(
+          pht(
+            'Storage has not been initialized yet, you must initialize '.
+            'storage before selectively applying patches.'));
+        return 1;
+      }
+
+      $legacy = $api->getLegacyPatches($this->patches);
+      if ($legacy || $no_quickstart || $init_only) {
+
+        // If we have legacy patches, we can't quickstart.
+
+        $api->createDatabase('meta_data');
+        $api->createTable(
+          'meta_data',
+          'patch_status',
+          array(
+            'patch VARCHAR(255) NOT NULL PRIMARY KEY COLLATE utf8_general_ci',
+            'applied INT UNSIGNED NOT NULL',
+          ));
+
+        foreach ($legacy as $patch) {
+          $api->markPatchApplied($patch);
+        }
+      } else {
+        echo pht('Loading quickstart template...')."\n";
+        $root = dirname(phutil_get_library_root('phabricator'));
+        $sql  = $root.'/resources/sql/quickstart.sql';
+        $api->applyPatchSQL($sql);
+      }
+    }
+
+    if ($init_only) {
+      echo pht('Storage initialized.')."\n";
+      return 0;
+    }
+
+    $applied = $api->getAppliedPatches();
+    $applied = array_fuse($applied);
+
+    $skip_mark = false;
+    if ($apply_only) {
+      if (isset($applied[$apply_only])) {
+
+        unset($applied[$apply_only]);
+        $skip_mark = true;
+
+        if (!$this->force && !$this->dryRun) {
+          echo phutil_console_wrap(
+            pht(
+              "Patch '%s' has already been applied. Are you sure you want ".
+              "to apply it again? This may put your storage in a state ".
+              "that the upgrade scripts can not automatically manage.",
+              $apply_only));
+          if (!phutil_console_confirm(pht('Apply patch again?'))) {
+            echo pht('Cancelled.')."\n";
+            return 1;
+          }
+        }
+      }
+    }
+
+    while (true) {
+      $applied_something = false;
+      foreach ($this->patches as $key => $patch) {
+        if (isset($applied[$key])) {
+          unset($this->patches[$key]);
+          continue;
+        }
+
+        if ($apply_only && $apply_only != $key) {
+          unset($this->patches[$key]);
+          continue;
+        }
+
+        $can_apply = true;
+        foreach ($patch->getAfter() as $after) {
+          if (empty($applied[$after])) {
+            if ($apply_only) {
+              echo pht(
+                "Unable to apply patch '%s' because it depends ".
+                "on patch '%s', which has not been applied.\n",
+                $apply_only,
+                $after);
+              return 1;
+            }
+            $can_apply = false;
+            break;
+          }
+        }
+
+        if (!$can_apply) {
+          continue;
+        }
+
+        $applied_something = true;
+
+        if ($this->dryRun) {
+          echo pht("DRYRUN: Would apply patch '%s'.", $key)."\n";
+        } else {
+          echo pht("Applying patch '%s'...", $key)."\n";
+
+          $t_begin = microtime(true);
+          $api->applyPatch($patch);
+          $t_end = microtime(true);
+
+          if (!$skip_mark) {
+            $api->markPatchApplied($key, ($t_end - $t_begin));
+          }
+        }
+
+        unset($this->patches[$key]);
+        $applied[$key] = true;
+      }
+
+      if (!$applied_something) {
+        if (count($this->patches)) {
+          throw new Exception(
+            pht(
+              'Some patches could not be applied: %s',
+              implode(', ', array_keys($this->patches))));
+        } else if (!$this->dryRun && !$apply_only) {
+          echo pht(
+            "Storage is up to date. Use '%s' for details.",
+            'storage status')."\n";
+        }
+        break;
+      }
+    }
+  }
+
   final protected function getBareHostAndPort($host) {
     // Split out port information, since the command-line client requires a
     // separate flag for the port.
@@ -680,4 +888,15 @@ abstract class PhabricatorStorageManagementWorkflow
     return array($bare_hostname, $port);
   }
 
+  /**
+   * Acquires a @{class:PhabricatorGlobalLock}.
+   *
+   * @return PhabricatorGlobalLock
+   */
+  final protected function lock() {
+    return PhabricatorGlobalLock::newLock(__CLASS__)
+      ->useSpecificConnection($this->getApi()->getConn(null))
+      ->lock();
+  }
+
 }

src/infrastructure/util/PhabricatorGlobalLock.php

@@ -62,6 +62,21 @@ final class PhabricatorGlobalLock extends PhutilLock {
     return $lock;
   }
 
+  /**
+   * Use a specific database connection for locking.
+   *
+   * By default, `PhabricatorGlobalLock` will lock on the "repository" database
+   * (somewhat arbitrarily). In most cases this is fine, but this method can
+   * be used to lock on a specific connection.
+   *
+   * @param  AphrontDatabaseConnection
+   * @return this
+   */
+  public function useSpecificConnection(AphrontDatabaseConnection $conn) {
+    $this->conn = $conn;
+    return $this;
+  }
+
 
 /* -(  Implementation  )----------------------------------------------------- */
 
@@ -86,13 +101,13 @@ final class PhabricatorGlobalLock extends PhutilLock {
       // NOTE: Using "force_new" to make sure each lock is on its own
       // connection.
       $conn = $dao->establishConnection('w', $force_new = true);
+    }
 
     // NOTE: Since MySQL will disconnect us if we're idle for too long, we set
     // the wait_timeout to an enormous value, to allow us to hold the
     // connection open indefinitely (or, at least, for 24 days).
     $max_allowed_timeout = 2147483;
     queryfx($conn, 'SET wait_timeout = %d', $max_allowed_timeout);
-    }
 
     $result = queryfx_one(
       $conn,