diff --git a/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php b/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php index e38b381cc5..af85f6bbec 100644 --- a/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php +++ b/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php @@ -1,288 +1,364 @@ setName('lease') ->setSynopsis(pht('Lease a resource.')) ->setArguments( array( array( 'name' => 'type', 'param' => 'resource_type', 'help' => pht('Resource type.'), ), array( 'name' => 'until', 'param' => 'time', 'help' => pht('Set lease expiration time.'), ), array( 'name' => 'attributes', 'param' => 'file', 'help' => pht( 'JSON file with lease attributes. Use "-" to read attributes '. 'from stdin.'), ), array( 'name' => 'count', 'param' => 'N', 'default' => 1, 'help' => pht('Lease a given number of identical resources.'), ), + array( + 'name' => 'blueprint', + 'param' => 'identifier', + 'repeat' => true, + 'help' => pht('Lease resources from a specific blueprint.'), + ), )); } public function execute(PhutilArgumentParser $args) { $viewer = $this->getViewer(); $resource_type = $args->getArg('type'); if (!phutil_nonempty_string($resource_type)) { throw new PhutilArgumentUsageException( pht( 'Specify a resource type with "--type".')); } $until = $args->getArg('until'); if (phutil_nonempty_string($until)) { $until = strtotime($until); if ($until <= 0) { throw new PhutilArgumentUsageException( pht( 'Unable to parse argument to "--until".')); } } $count = $args->getArgAsInteger('count'); if ($count < 1) { throw new PhutilArgumentUsageException( pht( 'Value provided to "--count" must be a nonzero, positive '. 'number.')); } $attributes_file = $args->getArg('attributes'); if (phutil_nonempty_string($attributes_file)) { if ($attributes_file == '-') { echo tsprintf( "%s\n", pht('Reading JSON attributes from stdin...')); $data = file_get_contents('php://stdin'); } else { $data = Filesystem::readFile($attributes_file); } $attributes = phutil_json_decode($data); } else { $attributes = array(); } + $filter_identifiers = $args->getArg('blueprint'); + if ($filter_identifiers) { + $filter_blueprints = $this->getBlueprintFilterMap($filter_identifiers); + } else { + $filter_blueprints = array(); + } + $blueprint_phids = null; $leases = array(); for ($idx = 0; $idx < $count; $idx++) { $lease = id(new DrydockLease()) ->setResourceType($resource_type); $drydock_phid = id(new PhabricatorDrydockApplication())->getPHID(); $lease->setAuthorizingPHID($drydock_phid); if ($attributes) { $lease->setAttributes($attributes); } if ($blueprint_phids === null) { - $blueprint_phids = $this->newAllowedBlueprintPHIDs($lease); + $blueprint_phids = $this->newAllowedBlueprintPHIDs( + $lease, + $filter_blueprints); } $lease->setAllowedBlueprintPHIDs($blueprint_phids); if ($until) { $lease->setUntil($until); } // If something fatals or the user interrupts the process (for example, // with "^C"), release the lease. We'll cancel this below, if the lease // actually activates. $lease->setReleaseOnDestruction(true); $leases[] = $lease; } // TODO: This would probably be better handled with PhutilSignalRouter, // but it currently doesn't route SIGINT. We're initializing it to setup // SIGTERM handling and make eventual migration easier. $router = PhutilSignalRouter::getRouter(); pcntl_signal(SIGINT, array($this, 'didReceiveInterrupt')); $t_start = microtime(true); echo tsprintf( "%s\n\n", pht('Leases queued for activation:')); foreach ($leases as $lease) { $lease->queueForActivation(); echo tsprintf( " __%s__\n", PhabricatorEnv::getProductionURI($lease->getURI())); } echo tsprintf( "\n%s\n\n", pht('Waiting for daemons to activate leases...')); foreach ($leases as $lease) { $this->waitUntilActive($lease); } // Now that we've survived activation and the lease is good, make it // durable. foreach ($leases as $lease) { $lease->setReleaseOnDestruction(false); } $t_end = microtime(true); echo tsprintf( "\n%s\n\n", pht( 'Activation complete. Leases are permanent until manually '. 'released with:')); foreach ($leases as $lease) { echo tsprintf( " %s\n", pht('$ ./bin/drydock release-lease --id %d', $lease->getID())); } echo tsprintf( "\n%s\n", pht( 'Leases activated in %sms.', new PhutilNumber((int)(($t_end - $t_start) * 1000)))); return 0; } public function didReceiveInterrupt($signo) { // Doing this makes us run destructors, particularly the "release on // destruction" trigger on the lease. exit(128 + $signo); } private function waitUntilActive(DrydockLease $lease) { $viewer = $this->getViewer(); $log_cursor = 0; $log_types = DrydockLogType::getAllLogTypes(); $is_active = false; while (!$is_active) { $lease->reload(); $pager = id(new AphrontCursorPagerView()) ->setBeforeID($log_cursor); // While we're waiting, show the user any logs which the daemons have // generated to give them some clue about what's going on. $logs = id(new DrydockLogQuery()) ->setViewer($viewer) ->withLeasePHIDs(array($lease->getPHID())) ->executeWithCursorPager($pager); if ($logs) { $logs = mpull($logs, null, 'getID'); ksort($logs); $log_cursor = last_key($logs); } foreach ($logs as $log) { $type_key = $log->getType(); if (isset($log_types[$type_key])) { $type_object = id(clone $log_types[$type_key]) ->setLog($log) ->setViewer($viewer); $log_data = $log->getData(); $type = $type_object->getLogTypeName(); $data = $type_object->renderLogForText($log_data); } else { $type = pht('Unknown ("%s")', $type_key); $data = null; } echo tsprintf( "(Lease #%d) <%s> %B\n", $lease->getID(), $type, $data); } $status = $lease->getStatus(); switch ($status) { case DrydockLeaseStatus::STATUS_ACTIVE: $is_active = true; break; case DrydockLeaseStatus::STATUS_RELEASED: throw new Exception(pht('Lease has already been released!')); case DrydockLeaseStatus::STATUS_DESTROYED: throw new Exception(pht('Lease has already been destroyed!')); case DrydockLeaseStatus::STATUS_BROKEN: throw new Exception(pht('Lease has been broken!')); case DrydockLeaseStatus::STATUS_PENDING: case DrydockLeaseStatus::STATUS_ACQUIRED: break; default: throw new Exception( pht( 'Lease has unknown status "%s".', $status)); } if ($is_active) { break; } else { sleep(1); } } } - private function newAllowedBlueprintPHIDs(DrydockLease $lease) { + private function getBlueprintFilterMap(array $identifiers) { + $viewer = $this->getViewer(); + + $query = id(new DrydockBlueprintQuery()) + ->setViewer($viewer) + ->withIdentifiers($identifiers); + + $blueprints = $query->execute(); + $blueprints = mpull($blueprints, null, 'getPHID'); + + $map = $query->getIdentifierMap(); + + $seen = array(); + foreach ($identifiers as $identifier) { + if (!isset($map[$identifier])) { + throw new PhutilArgumentUsageException( + pht( + 'Blueprint "%s" could not be loaded. Try a blueprint ID or '. + 'PHID.', + $identifier)); + } + + $blueprint = $map[$identifier]; + + $blueprint_phid = $blueprint->getPHID(); + if (isset($seen[$blueprint_phid])) { + throw new PhutilArgumentUsageException( + pht( + 'Blueprint "%s" is specified more than once (as "%s" and "%s").', + $blueprint->getBlueprintName(), + $seen[$blueprint_phid], + $identifier)); + } + + $seen[$blueprint_phid] = true; + } + + return mpull($map, null, 'getPHID'); + } + + private function newAllowedBlueprintPHIDs( + DrydockLease $lease, + array $filter_blueprints) { + assert_instances_of($filter_blueprints, 'DrydockBlueprint'); + $viewer = $this->getViewer(); $impls = DrydockBlueprintImplementation::getAllForAllocatingLease($lease); if (!$impls) { throw new PhutilArgumentUsageException( pht( 'No known blueprint class can ever allocate the specified '. 'lease. Check that the resource type is spelled correctly.')); } $classes = array_keys($impls); $blueprints = id(new DrydockBlueprintQuery()) ->setViewer($viewer) ->withBlueprintClasses($classes) ->withDisabled(false) ->execute(); if (!$blueprints) { throw new PhutilArgumentUsageException( pht( 'No enabled blueprints exist with a blueprint class that can '. 'plausibly allocate resources to satisfy the requested lease.')); } $phids = mpull($blueprints, 'getPHID'); + if ($filter_blueprints) { + $allowed_map = array_fuse($phids); + $filter_map = mpull($filter_blueprints, null, 'getPHID'); + + foreach ($filter_map as $filter_phid => $blueprint) { + if (!isset($allowed_map[$filter_phid])) { + throw new PhutilArgumentUsageException( + pht( + 'Specified blueprint "%s" is not capable of satisfying the '. + 'configured lease.', + $blueprint->getBlueprintName())); + } + } + + $phids = mpull($filter_blueprints, 'getPHID'); + } + return $phids; } } diff --git a/src/applications/drydock/query/DrydockBlueprintQuery.php b/src/applications/drydock/query/DrydockBlueprintQuery.php index 6c92927bb8..a5df734e19 100644 --- a/src/applications/drydock/query/DrydockBlueprintQuery.php +++ b/src/applications/drydock/query/DrydockBlueprintQuery.php @@ -1,144 +1,241 @@ ids = $ids; return $this; } public function withPHIDs(array $phids) { $this->phids = $phids; return $this; } public function withBlueprintClasses(array $classes) { $this->blueprintClasses = $classes; return $this; } public function withDatasourceQuery($query) { $this->datasourceQuery = $query; return $this; } public function withDisabled($disabled) { $this->disabled = $disabled; return $this; } public function withAuthorizedPHIDs(array $phids) { $this->authorizedPHIDs = $phids; return $this; } public function withNameNgrams($ngrams) { return $this->withNgramsConstraint( new DrydockBlueprintNameNgrams(), $ngrams); } + public function withIdentifiers(array $identifiers) { + if (!$identifiers) { + throw new Exception( + pht( + 'Can not issue a query with an empty identifier list.')); + } + + $this->identifiers = $identifiers; + + $ids = array(); + $phids = array(); + + foreach ($identifiers as $identifier) { + if (ctype_digit($identifier)) { + $ids[] = $identifier; + } else { + $phids[] = $identifier; + } + } + + $this->identifierIDs = $ids; + $this->identifierPHIDs = $phids; + + return $this; + } + + public function getIdentifierMap() { + if ($this->identifierMap === null) { + throw new Exception( + pht( + 'Execute a query with identifiers before getting the '. + 'identifier map.')); + } + + return $this->identifierMap; + } + public function newResultObject() { return new DrydockBlueprint(); } protected function getPrimaryTableAlias() { return 'blueprint'; } protected function loadPage() { return $this->loadStandardPage($this->newResultObject()); } + protected function willExecute() { + if ($this->identifiers) { + $this->identifierMap = array(); + } else { + $this->identifierMap = null; + } + } + protected function willFilterPage(array $blueprints) { $impls = DrydockBlueprintImplementation::getAllBlueprintImplementations(); foreach ($blueprints as $key => $blueprint) { $impl = idx($impls, $blueprint->getClassName()); if (!$impl) { $this->didRejectResult($blueprint); unset($blueprints[$key]); continue; } $impl = clone $impl; $blueprint->attachImplementation($impl); } + if ($this->identifiers) { + $id_map = mpull($blueprints, null, 'getID'); + $phid_map = mpull($blueprints, null, 'getPHID'); + + $map = $this->identifierMap; + + foreach ($this->identifierIDs as $id) { + if (isset($id_map[$id])) { + $map[$id] = $id_map[$id]; + } + } + + foreach ($this->identifierPHIDs as $phid) { + if (isset($phid_map[$phid])) { + $map[$phid] = $phid_map[$phid]; + } + } + + // Just for consistency, reorder the map to match input order. + $map = array_select_keys($map, $this->identifiers); + + $this->identifierMap = $map; + } + return $blueprints; } protected function buildWhereClauseParts(AphrontDatabaseConnection $conn) { $where = parent::buildWhereClauseParts($conn); if ($this->ids !== null) { $where[] = qsprintf( $conn, 'blueprint.id IN (%Ld)', $this->ids); } if ($this->phids !== null) { $where[] = qsprintf( $conn, 'blueprint.phid IN (%Ls)', $this->phids); } if ($this->datasourceQuery !== null) { $where[] = qsprintf( $conn, 'blueprint.blueprintName LIKE %>', $this->datasourceQuery); } if ($this->blueprintClasses !== null) { $where[] = qsprintf( $conn, 'blueprint.className IN (%Ls)', $this->blueprintClasses); } if ($this->disabled !== null) { $where[] = qsprintf( $conn, 'blueprint.isDisabled = %d', (int)$this->disabled); } + if ($this->identifiers !== null) { + $parts = array(); + + if ($this->identifierIDs) { + $parts[] = qsprintf( + $conn, + 'blueprint.id IN (%Ld)', + $this->identifierIDs); + } + + if ($this->identifierPHIDs) { + $parts[] = qsprintf( + $conn, + 'blueprint.phid IN (%Ls)', + $this->identifierPHIDs); + } + + $where[] = qsprintf( + $conn, + '%LO', + $parts); + } + return $where; } protected function shouldGroupQueryResultRows() { if ($this->authorizedPHIDs !== null) { return true; } return parent::shouldGroupQueryResultRows(); } protected function buildJoinClauseParts(AphrontDatabaseConnection $conn) { $joins = parent::buildJoinClauseParts($conn); if ($this->authorizedPHIDs !== null) { $joins[] = qsprintf( $conn, 'JOIN %T authorization ON authorization.blueprintPHID = blueprint.phid AND authorization.objectPHID IN (%Ls) AND authorization.objectAuthorizationState = %s AND authorization.blueprintAuthorizationState = %s', id(new DrydockAuthorization())->getTableName(), $this->authorizedPHIDs, DrydockAuthorization::OBJECTAUTH_ACTIVE, DrydockAuthorization::BLUEPRINTAUTH_AUTHORIZED); } return $joins; } }