com.staros.worker.ServiceWorkerGroup 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.worker;

import com.staros.util.Constant;
import com.staros.util.LogUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.HashMap;
import java.util.Map;

/**
 * Service Worker Group manages all workers in one service, all
 * workers are grouped by WorkerGroup in Service Worker Group.
 */
public class ServiceWorkerGroup {
    private static final Logger LOG = LogManager.getLogger(ServiceWorkerGroup.class);

    private final String serviceId;
    private Map workerGroups; // 
    // TODO:
    //  1. update defaultWorkerGroupId when support deleting workerGroup operation.
    //  2. Cancel WorkerGroup-0 default workerGroup implementation
    private WorkerGroup defaultWorkerGroup;

    public ServiceWorkerGroup(String serviceId) {
        // TODO: load from storage
        this.serviceId = serviceId;
        this.workerGroups = new HashMap<>();
        this.defaultWorkerGroup = new WorkerGroup(serviceId, Constant.DEFAULT_ID);

        this.workerGroups.put(defaultWorkerGroup.getGroupId(), defaultWorkerGroup);
    }

    public void createGroup(long groupId) {
        if (workerGroups.containsKey(groupId)) { // should not happen
            LogUtils.fatal(LOG, "group id {} already exists when create group in service {}!", groupId, serviceId);
        }

        WorkerGroup group = new WorkerGroup(serviceId, groupId);
        workerGroups.put(groupId, group);
        if (defaultWorkerGroup.getGroupId() == Constant.DEFAULT_ID) {
            defaultWorkerGroup = group;
        }
    }

    public void addWorker(Worker worker) {
        checkWorker(worker);
        getWorkerGroup(worker.getGroupId()).addWorker(worker);
    }

    public void removeWorker(Worker worker) {
        checkWorker(worker);
        getWorkerGroup(worker.getGroupId()).removeWorker(worker);
    }

    public void updateWorker(Worker worker) {
        checkWorker(worker);
        getWorkerGroup(worker.getGroupId()).updateWorker(worker);
    }

    public WorkerGroup getDefaultWorkerGroup() {
        return defaultWorkerGroup;
    }

    public WorkerGroup getWorkerGroup(long groupId) {
        return workerGroups.get(groupId);
    }

    // crash on purpose if check failed
    private void checkWorker(Worker worker) {
        if (!worker.getServiceId().equals(serviceId)) {
            LogUtils.fatal(LOG, "worker's service id {} not match service id {} in service worker group.",
                    worker.getServiceId(), serviceId);
        }

        WorkerGroup workerGroup = workerGroups.get(worker.getGroupId());
        if (workerGroup == null) {
            LogUtils.fatal(LOG, "worker's group id {} not exist in service worker group.",
                    worker.getGroupId());
        }
    }

    public boolean existGroup(long groupId) {
        return workerGroups.containsKey(groupId);
    }
}