com.staros.schedule.ScheduleScorer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of starmanager Show documentation
There is a newer version: 3.4-rc2
// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


package com.staros.schedule;

import com.google.common.collect.ImmutableMap;
import com.staros.proto.PlacementPolicy;
import com.staros.schedule.select.Selector;
import com.staros.worker.Worker;
import com.staros.worker.WorkerManager;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ScheduleScorer {
    /**
     * WEIGHT defines that the higher weight, the more preferable the worker will be chosen.
     * TODO: make these numbers configurable.
     */
    // A REVERSE weight, the larger the weight given, the less preferred the worker will be selected.
    private static final long workerShardNumberWeight = 1;
    // worker with existing pack replica are greatly preferred.
    private static final long shardGroupPackWeight = 10000;
    // worker with existing spread replica are not preferred
    private static final long shardGroupSpreadWeight = -100;

    private static final Map shardGroupWeight =
            ImmutableMap.builder()
                    .put(PlacementPolicy.PACK, shardGroupPackWeight)
                    .put(PlacementPolicy.SPREAD, shardGroupSpreadWeight)
                    .build();

    private final Map scores;

    public ScheduleScorer(Collection targetWorkers) {
        this.scores = new HashMap<>(targetWorkers.size());
        targetWorkers.forEach(x -> this.scores.put(x, 0.0));
    }

    /**
     * Allow query worker manager about the worker info and change the score accordingly.
     *
     * @param manager WorkerManager to retrieve the worker status
     */
    public void apply(WorkerManager manager) {
        // TODO: take more metrics into consideration
        this.scores.replaceAll((id, val) -> {
            Worker worker = manager.getWorker(id);
            return worker == null ? val : val - Math.log(worker.getNumOfShards() * workerShardNumberWeight + 1);
        });
    }

    public void apply(PlacementPolicy policy, Collection workerIds) {
        if (!shardGroupWeight.containsKey(policy)) {
            return;
        }
        long weight = shardGroupWeight.get(policy);
        workerIds.forEach(key -> scores.computeIfPresent(key, (k, val) -> val + weight));
    }

    public boolean isEmpty() {
        return scores.isEmpty();
    }

    public void remove(long workerId) {
        scores.remove(workerId);
    }

    public Map getScores() {
        return Collections.unmodifiableMap(scores);
    }

    public List selectHighEnd(Selector selector, int nSelect) {
        return selector.selectHighEnd(scores, nSelect);
    }

    public List selectLowEnd(Selector selector, int nSelect) {
        return selector.selectLowEnd(scores, nSelect);
    }
}