cn.ponfee.disjob.registry.discovery.WorkerDiscovery Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2022-2024 Ponfee (http://www.ponfee.cn/)
 *
 * 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 cn.ponfee.disjob.registry.discovery;

import cn.ponfee.disjob.core.base.*;
import cn.ponfee.disjob.core.dto.worker.SubscribeSupervisorChangedParam;
import cn.ponfee.disjob.registry.rpc.DestinationServerRestProxy;
import cn.ponfee.disjob.registry.rpc.DestinationServerRestProxy.DestinationServerClient;
import com.google.common.collect.ImmutableCollection;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.util.Assert;
import org.springframework.web.client.RestTemplate;

import java.util.*;
import java.util.stream.Collectors;

/**
 * Worker discovery.
 *
 * @author Ponfee
 */
public final class WorkerDiscovery extends ServerDiscovery {

    private final DestinationServerClient workerRpcClient;

    /**
     * Map>
     */
    private volatile ImmutableMap> groupedWorkers = ImmutableMap.of();

    WorkerDiscovery(RestTemplate restTemplate) {
        this.workerRpcClient = DestinationServerRestProxy.create(
            WorkerRpcService.class,
            null,
            null,
            worker -> Supervisor.local().getWorkerContextPath(worker.getGroup()),
            restTemplate,
            RetryProperties.none()
        );
    }

    @Override
    public synchronized void refreshServers(List discoveredWorkers) {
        if (CollectionUtils.isEmpty(discoveredWorkers)) {
            this.groupedWorkers = ImmutableMap.of();
        } else {
            this.groupedWorkers = discoveredWorkers.stream()
                .collect(Collectors.groupingBy(Worker::getGroup))
                .entrySet()
                .stream()
                .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, e -> toSortedImmutableList(e.getValue())));
        }
    }

    @Override
    public synchronized void updateServers(RegistryEventType eventType, Worker worker) {
        if (eventType.isRegister() && isAlive(worker)) {
            return;
        }
        if (eventType.isDeregister() && !isAlive(worker)) {
            return;
        }
        ImmutableMap.Builder> builder = ImmutableMap.builder();
        groupedWorkers.forEach((k, v) -> builder.put(k, worker.equalsGroup(k) ? mergeServers(v, eventType, worker) : v));
        if (!groupedWorkers.containsKey(worker.getGroup())) {
            builder.put(worker.getGroup(), ImmutableList.of(worker));
        }
        this.groupedWorkers = builder.build();
    }

    @Override
    public List getServers(String group) {
        Assert.hasText(group, "Get discovery worker group cannot null.");
        List workers = groupedWorkers.get(group);
        return workers == null ? Collections.emptyList() : workers;
    }

    @Override
    public boolean hasServers() {
        return !groupedWorkers.isEmpty();
    }

    @Override
    public boolean isAlive(Worker worker) {
        List workers = groupedWorkers.get(worker.getGroup());
        return workers != null && Collections.binarySearch(workers, worker) > -1;
    }

    // ----------------------------------------------------------------default package methods

    @Override
    List getServers() {
        ImmutableCollection> values = groupedWorkers.values();
        List list = new ArrayList<>(values.stream().mapToInt(AbstractCollection::size).sum());
        values.forEach(list::addAll);
        return list;
    }

    @Override
    void notifyServer(Worker worker, RegistryEventType eventType, Supervisor supervisor) {
        /*
        if (worker.matches(Worker.local())) {
            return;
        }
        */
        try {
            String authToken = Supervisor.local().createSupervisorAuthenticationToken(worker.getGroup());
            SubscribeSupervisorChangedParam param = SubscribeSupervisorChangedParam.of(authToken, eventType, supervisor);
            workerRpcClient.invoke(worker, client -> client.subscribeSupervisorChanged(param));
        } catch (Throwable t) {
            log.error("Notify server error: {}, {}", worker, t.getMessage());
        }
    }

}