summaryrefslogtreecommitdiff
path: root/jstests/libs/override_methods/multiversion_override_balancer_control.js
blob: b24d209a7d6a28ad8798650f4f9930458ebfbf3a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
/**
 * This test overrides the ShardingTest start/stopBalancer methods to go directly against the config
 * server. The reason is that tests with 3.2 mongos and 3.4 mongod are failing because in that
 * intermediate state neither the new shell is able to stop the balancer (missing balancerStart/Stop
 * commands on mongos), nor the old shell works (because the balancer lock session id does not
 * change since the balancer lock is held permanently).
 */
(function() {
    'use strict';

    // Preserve the original ShardingTest constructor, because we are overriding it
    var originalShardingTest = ShardingTest;

    ShardingTest = function ShardingTestWithMongoSOnLatestVersionForBalancerControl() {
        // The ShardingTest constructor uses the stopBalancer method if the 'enableBalancer' setting
        // is not present. For this reason, temporary patch it to be 'enabled' and then call
        // stopBalancer afterwards.
        if (!arguments[0].other) {
            arguments[0].other = {};
        }

        // Do not start the mongos server for controlling the balancer if the config server is 3.2.
        var skipControlMongoS = false;

        var configOptions = arguments[0].other.configOptions;
        if (configOptions && configOptions.binVersion == "last-stable") {
            skipControlMongoS = true;
        }

        // Check if config server options were specified as an array instead.
        if (arguments[0].config && Array.isArray(arguments[0].config) &&
            arguments[0].config[0].binVersion == "last-stable") {
            skipControlMongoS = true;
        }

        if (skipControlMongoS) {
            // Construct the original object and skip creating the separate mongos for controlling
            // the balancer.
            originalShardingTest.apply(this, arguments);
            return;
        }

        var originalEnableBalancer = arguments[0].other.enableBalancer;
        if (!originalEnableBalancer) {
            arguments[0].other.enableBalancer = true;
        }

        // Construct the original object
        originalShardingTest.apply(this, arguments);

        // Start one mongos on the side to be used for balancer control
        var controlMongoSStartupOptions = {configdb: this._configDB, binVersion: 'latest'};

        if (arguments[0].other.keyFile) {
            controlMongoSStartupOptions.keyFile = arguments[0].other.keyFile;
        }

        print('Starting balancer control mongos instance ...');
        var controlMongoSConn = MongoRunner.runMongos(controlMongoSStartupOptions);

        // Override the start/stopBalancer methods

        var originalStartBalancer = this.startBalancer;
        this.startBalancer = function(timeoutMs, interval) {
            timeoutMs = timeoutMs || 60000;

            var result;
            var fn = controlMongoSConn.adminCommand.bind(controlMongoSConn,
                                                         {balancerStart: 1, maxTimeMS: timeoutMs});
            if (controlMongoSStartupOptions.keyFile) {
                result =
                    authutil.asCluster(controlMongoSConn, controlMongoSStartupOptions.keyFile, fn);
            } else {
                result = fn();
            }

            // Back off to the legacy control if the command is not supported
            if (result.code === ErrorCodes.CommandNotFound) {
                return originalStartBalancer.apply(this, [timeoutMs, interval]);
            }

            return assert.commandWorked(result);
        };

        var originalStopBalancer = this.stopBalancer;
        this.stopBalancer = function(timeoutMs, interval) {
            timeoutMs = timeoutMs || 60000;

            var result;
            var fn = controlMongoSConn.adminCommand.bind(controlMongoSConn,
                                                         {balancerStop: 1, maxTimeMS: timeoutMs});
            if (controlMongoSStartupOptions.keyFile) {
                result =
                    authutil.asCluster(controlMongoSConn, controlMongoSStartupOptions.keyFile, fn);
            } else {
                result = fn();
            }

            // Back off to the legacy control if the command is not supported
            if (result.code === ErrorCodes.CommandNotFound) {
                return originalStopBalancer.apply(this, [timeoutMs, interval]);
            }

            return assert.commandWorked(result);
        };

        var originalAwaitBalancerRound = this.awaitBalancerRound;
        this.awaitBalancerRound = function(timeoutMs) {
            timeoutMs = timeoutMs || 60000;

            var fn = originalAwaitBalancerRound.bind(this, timeoutMs, controlMongoSConn);
            if (controlMongoSStartupOptions.keyFile) {
                authutil.asCluster(controlMongoSConn, controlMongoSStartupOptions.keyFile, fn);
            } else {
                fn();
            }
        };

        // Override the stop method to also stop the control mongos
        var originalStop = this.stop;
        this.stop = function() {
            MongoRunner.stopMongos(controlMongoSConn);
            originalStop.apply(this);
        };

        // Start/stop the balancer as requested
        if (!originalEnableBalancer) {
            arguments[0].other.enableBalancer = originalEnableBalancer;
            this.stopBalancer();
        }
    };

    Object.extend(ShardingTest, originalShardingTest);

})();