com.tencent.angel.master.client.ClientManager Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of angel-ps-core Show documentation
The newest version!
/*
 * Tencent is pleased to support the open source community by making Angel available.
 *
 * Copyright (C) 2017-2018 THL A29 Limited, a Tencent company. 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://opensource.org/licenses/Apache-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.tencent.angel.master.client;

import com.tencent.angel.common.location.Location;
import com.tencent.angel.conf.AngelConf;
import com.tencent.angel.master.app.AMContext;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Angel client manager
 */
public class ClientManager {
  private static final Log LOG = LogFactory.getLog(ClientManager.class);
  private final AMContext context;

  /**
   * Client id to last heartbeat timestamp map
   */
  private final ConcurrentHashMap clientToLastHBTsMap;

  /**
   * Client id to client location map
   */
  private final ConcurrentHashMap clientIdToLocMap;

  /**
   * Client timeout time in milliseconds
   */
  private final long clientTimeoutMS;

  /**
   * Client id generator
   */
  private final AtomicInteger idGen = new AtomicInteger(0);

  public ClientManager(AMContext context) {
    this.context = context;
    clientToLastHBTsMap = new ConcurrentHashMap<>();
    clientIdToLocMap = new ConcurrentHashMap<>();
    clientTimeoutMS = context.getConf()
      .getLong(AngelConf.ANGEL_CLIENT_HEARTBEAT_INTERVAL_TIMEOUT_MS,
        AngelConf.DEFAULT_ANGEL_CLIENT_HEARTBEAT_INTERVAL_TIMEOUT_MS);
  }

  /**
   * Check if some client heartbeat timeout, if all client timeout, just kill the application
   */
  public void checkHBTimeOut() {
    Iterator> clientIt = clientToLastHBTsMap.entrySet().iterator();
    boolean isTimeOut = false;
    Map.Entry clientEntry;
    long currentTs = System.currentTimeMillis();
    while (clientIt.hasNext()) {
      clientEntry = clientIt.next();
      if (currentTs - clientEntry.getValue() > clientTimeoutMS) {
        LOG.error("Client " + clientEntry.getKey() + " heartbeat timeout");
        clientIt.remove();
        isTimeOut = true;
      }
    }

    if (isTimeOut && clientToLastHBTsMap.isEmpty()) {
      LOG.error("All client timeout, just exit the application");
      context.getMasterService().stop(1);
    }
  }

  /**
   * Is a client alive
   *
   * @param clientId client id
   * @return true means alive
   */
  public boolean isAlive(int clientId) {
    return clientToLastHBTsMap.containsKey(clientId);
  }

  /**
   * Update the heartbeat timestamp for a client
   *
   * @param clientId client id
   */
  public void alive(int clientId) {
    clientToLastHBTsMap.put(clientId, System.currentTimeMillis());
  }

  /**
   * Client register
   *
   * @param clientId client id
   */
  public void register(int clientId) {
    clientToLastHBTsMap.put(clientId, System.currentTimeMillis());
  }

  /**
   * Client unregister
   *
   * @param clientId client id
   */
  public void unRegister(int clientId) {
    clientToLastHBTsMap.remove(clientId);
  }

  /**
   * Generate a new client id
   *
   * @return a new client id
   */
  public int getId() {
    return idGen.incrementAndGet();
  }
}