LBFactoryMulti.php

Go to the documentation of this file.

<?php
namespace Wikimedia\Rdbms;

use InvalidArgumentException;

class LBFactoryMulti extends LBFactory {
    private $sectionsByDB;
    private $sectionLoads;
    private $serverTemplate;

    private $groupLoadsBySection = [];
    private $groupLoadsByDB = [];
    private $hostsByName = [];
    private $externalLoads = [];
    private $externalTemplateOverrides;
    private $templateOverridesByServer;
    private $templateOverridesBySection;
    private $templateOverridesByCluster;
    private $masterTemplateOverrides;
    private $readOnlyBySection = [];

    private $mainLBs = [];
    private $extLBs = [];
    private $loadMonitorClass = 'LoadMonitor';
    private $lastDomain;
    private $lastSection;

    public function __construct( array $conf ) {
        parent::__construct( $conf );

        $required = [ 'sectionsByDB', 'sectionLoads', 'serverTemplate' ];
        $optional = [ 'groupLoadsBySection', 'groupLoadsByDB', 'hostsByName',
            'externalLoads', 'externalTemplateOverrides', 'templateOverridesByServer',
            'templateOverridesByCluster', 'templateOverridesBySection', 'masterTemplateOverrides',
            'readOnlyBySection', 'loadMonitorClass' ];

        foreach ( $required as $key ) {
            if ( !isset( $conf[$key] ) ) {
                throw new InvalidArgumentException( __CLASS__ . ": $key is required." );
            }
            $this->$key = $conf[$key];
        }

        foreach ( $optional as $key ) {
            if ( isset( $conf[$key] ) ) {
                $this->$key = $conf[$key];
            }
        }
    }

    private function getSectionForDomain( $domain = false ) {
        if ( $this->lastDomain === $domain ) {
            return $this->lastSection;
        }

        $database = $this->getDatabaseFromDomain( $domain );
        $section = $this->sectionsByDB[$database] ?? 'DEFAULT';
        $this->lastSection = $section;
        $this->lastDomain = $domain;

        return $section;
    }

    public function newMainLB( $domain = false ) {
        $database = $this->getDatabaseFromDomain( $domain );
        $section = $this->getSectionForDomain( $domain );
        $groupLoads = $this->groupLoadsByDB[$database] ?? [];

        if ( isset( $this->groupLoadsBySection[$section] ) ) {
            $groupLoads = array_merge_recursive(
                $groupLoads, $this->groupLoadsBySection[$section] );
        }

        $readOnlyReason = $this->readOnlyReason;
        // Use the LB-specific read-only reason if everything isn't already read-only
        if ( $readOnlyReason === false && isset( $this->readOnlyBySection[$section] ) ) {
            $readOnlyReason = $this->readOnlyBySection[$section];
        }

        $template = $this->serverTemplate;
        if ( isset( $this->templateOverridesBySection[$section] ) ) {
            $template = $this->templateOverridesBySection[$section] + $template;
        }

        return $this->newLoadBalancer(
            $template,
            $this->sectionLoads[$section],
            $groupLoads,
            $readOnlyReason
        );
    }

    public function getMainLB( $domain = false ) {
        $section = $this->getSectionForDomain( $domain );
        if ( !isset( $this->mainLBs[$section] ) ) {
            $this->mainLBs[$section] = $this->newMainLB( $domain );
        }

        return $this->mainLBs[$section];
    }

    public function newExternalLB( $cluster ) {
        if ( !isset( $this->externalLoads[$cluster] ) ) {
            throw new InvalidArgumentException( __METHOD__ . ": Unknown cluster \"$cluster\"" );
        }
        $template = $this->serverTemplate;
        if ( $this->externalTemplateOverrides ) {
            $template = $this->externalTemplateOverrides + $template;
        }
        if ( isset( $this->templateOverridesByCluster[$cluster] ) ) {
            $template = $this->templateOverridesByCluster[$cluster] + $template;
        }

        return $this->newLoadBalancer(
            $template,
            $this->externalLoads[$cluster],
            [],
            $this->readOnlyReason
        );
    }

    public function getExternalLB( $cluster ) {
        if ( !isset( $this->extLBs[$cluster] ) ) {
            $this->extLBs[$cluster] = $this->newExternalLB( $cluster );
        }

        return $this->extLBs[$cluster];
    }

    public function getAllMainLBs() {
        $lbs = [];
        foreach ( $this->sectionsByDB as $db => $section ) {
            if ( !isset( $lbs[$section] ) ) {
                $lbs[$section] = $this->getMainLB( $db );
            }
        }

        return $lbs;
    }

    public function getAllExternalLBs() {
        $lbs = [];
        foreach ( $this->externalLoads as $cluster => $unused ) {
            $lbs[$cluster] = $this->getExternalLB( $cluster );
        }

        return $lbs;
    }

    private function newLoadBalancer( $template, $loads, $groupLoads, $readOnlyReason ) {
        $lb = new LoadBalancer( array_merge(
            $this->baseLoadBalancerParams(),
            [
                'servers' => $this->makeServerArray( $template, $loads, $groupLoads ),
                'loadMonitor' => [ 'class' => $this->loadMonitorClass ],
                'readOnlyReason' => $readOnlyReason
            ]
        ) );
        $this->initLoadBalancer( $lb );

        return $lb;
    }

    private function makeServerArray( $template, $loads, $groupLoads ) {
        $servers = [];
        $master = true;
        $groupLoadsByServer = $this->reindexGroupLoads( $groupLoads );
        foreach ( $groupLoadsByServer as $server => $stuff ) {
            if ( !isset( $loads[$server] ) ) {
                $loads[$server] = 0;
            }
        }
        foreach ( $loads as $serverName => $load ) {
            $serverInfo = $template;
            if ( $master ) {
                $serverInfo['master'] = true;
                if ( $this->masterTemplateOverrides ) {
                    $serverInfo = $this->masterTemplateOverrides + $serverInfo;
                }
                $master = false;
            } else {
                $serverInfo['replica'] = true;
            }
            if ( isset( $this->templateOverridesByServer[$serverName] ) ) {
                $serverInfo = $this->templateOverridesByServer[$serverName] + $serverInfo;
            }
            if ( isset( $groupLoadsByServer[$serverName] ) ) {
                $serverInfo['groupLoads'] = $groupLoadsByServer[$serverName];
            }
            $serverInfo['host'] = $this->hostsByName[$serverName] ?? $serverName;
            $serverInfo['hostName'] = $serverName;
            $serverInfo['load'] = $load;
            $serverInfo += [ 'flags' => IDatabase::DBO_DEFAULT ];

            $servers[] = $serverInfo;
        }

        return $servers;
    }

    private function reindexGroupLoads( $groupLoads ) {
        $reindexed = [];
        foreach ( $groupLoads as $group => $loads ) {
            foreach ( $loads as $server => $load ) {
                $reindexed[$server][$group] = $load;
            }
        }

        return $reindexed;
    }

    private function getDatabaseFromDomain( $domain = false ) {
        return ( $domain === false )
            ? $this->localDomain->getDatabase()
            : DatabaseDomain::newFromId( $domain )->getDatabase();
    }

    public function forEachLB( $callback, array $params = [] ) {
        foreach ( $this->mainLBs as $lb ) {
            $callback( $lb, ...$params );
        }
        foreach ( $this->extLBs as $lb ) {
            $callback( $lb, ...$params );
        }
    }
}