content: /*
   *
   * Copyright 1999-2021 Alibaba Group Holding Ltd.
   *
   * 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
   *
   *      http://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.alibaba.nacos.core.remote;
  
  import com.alibaba.nacos.api.common.Constants;
  import com.alibaba.nacos.api.remote.RemoteConstants;
  import com.alibaba.nacos.api.remote.RpcScheduledExecutor;
  import com.alibaba.nacos.api.remote.request.ConnectResetRequest;
  import com.alibaba.nacos.common.remote.exception.ConnectionAlreadyClosedException;
  import com.alibaba.nacos.common.spi.NacosServiceLoader;
  import com.alibaba.nacos.common.utils.StringUtils;
  import com.alibaba.nacos.core.monitor.MetricsMonitor;
  import com.alibaba.nacos.plugin.control.ControlManagerCenter;
  import com.alibaba.nacos.plugin.control.Loggers;
  import com.alibaba.nacos.plugin.control.configs.ControlConfigs;
  import com.alibaba.nacos.plugin.control.connection.request.ConnectionCheckRequest;
  import com.alibaba.nacos.plugin.control.connection.response.ConnectionCheckResponse;
  import com.alibaba.nacos.plugin.control.connection.rule.ConnectionControlRule;
  import com.alibaba.nacos.sys.env.EnvUtil;
  import org.slf4j.Logger;
  import org.springframework.stereotype.Service;
  
  import javax.annotation.PostConstruct;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  
  /**
   * connect manager.
   *
   * @author liuzunfei
   * @version $Id: ConnectionManager.java, v 0.1 2020年07月13日 7:07 PM liuzunfei Exp $
   */
  @Service
  public class ConnectionManager {
      
      private static final Logger LOGGER = com.alibaba.nacos.plugin.control.Loggers.CONNECTION;
      
      private Map<String, AtomicInteger> connectionForClientIp = new ConcurrentHashMap<>(16);
      
      Map<String, Connection> connections = new ConcurrentHashMap<>();
      
      private RuntimeConnectionEjector runtimeConnectionEjector;
      
      private ClientConnectionEventListenerRegistry clientConnectionEventListenerRegistry;
      
      public ConnectionManager(ClientConnectionEventListenerRegistry clientConnectionEventListenerRegistry) {
          this.clientConnectionEventListenerRegistry = clientConnectionEventListenerRegistry;
      }
      
      /**
       * if monitor detail.
       *
       * @param clientIp clientIp.
       * @return
       */
      public boolean traced(String clientIp) {
          ConnectionControlRule connectionControlRule = ControlManagerCenter.getInstance().getConnectionControlManager()
                  .getConnectionLimitRule();
          return connectionControlRule != null && connectionControlRule.getMonitorIpList() != null
                  && connectionControlRule.getMonitorIpList().contains(clientIp);
      }
      
      /**
       * check connection id is valid.
       *
       * @param connectionId connectionId to be check.
       * @return is valid or not.
       */
      public boolean checkValid(String connectionId) {
          return connections.containsKey(connectionId);
      }
      
      /**
       * register a new connect.
       *
       * @param connectionId connectionId
       * @param connection   connection
       */
      public synchronized boolean register(String connectionId, Connection connection) {
          
          if (connection.isConnected()) {
              String clientIp = connection.getMetaInfo().clientIp;
              if (connections.containsKey(connectionId)) {
                  return true;
              }
              if (checkLimit(connection)) {
                  return false;
              }
              if (traced(clientIp)) {
                  connection.setTraced(true);
              }
              connections.put(connectionId, connection);
              connectionForClientIp.computeIfAbsent(clientIp, k -> new AtomicInteger(0)).getAndIncrement();
              
              clientConnectionEventListenerRegistry.notifyClientConnected(connection);
              
              LOGGER.info("new connection registered successfully, connectionId = {},connection={} ", connectionId,
                      connection);
              return true;
              
          }
          return false;
          
      }
      
      private boolean checkLimit(Connection connection) {
          if (connection.getMetaInfo().isClusterSource()) {
              return false;
          }
          ConnectionMeta metaInfo = connection.getMetaInfo();
          ConnectionCheckRequest connectionCheckRequest = new ConnectionCheckRequest(metaInfo.getClientIp(),
                  metaInfo.getAppName(), metaInfo.getLabel(RemoteConstants.LABEL_SOURCE));
          connectionCheckRequest.setLabels(connection.getLabels());
          ConnectionCheckResponse checkResponse = ControlManagerCenter.getInstance().getConnectionControlManager()
                  .check(connectionCheckRequest);
          return !checkResponse.isSuccess();
      }
      
      /**
       * unregister a connection .
       *
       * @param connectionId connectionId.
       */
      public synchronized void unregister(String connectionId) {
          Connection remove = this.connections.remove(connectionId);
          if (remove != null) {
              String clientIp = remove.getMetaInfo().clientIp;
              AtomicInteger atomicInteger = connectionForClientIp.get(clientIp);
              if (atomicInteger != null) {
                  int count = atomicInteger.decrementAndGet();
                  if (count <= 0) {
                      connectionForClientIp.remove(clientIp);
                  }
              }
              remove.close();
              LOGGER.info("[{}]Connection unregistered successfully. ", connectionId);
              clientConnectionEventListenerRegistry.notifyClientDisConnected(remove);
          }
      }
      
      /**
       * get by connection id.
       *
       * @param connectionId connection id.
       * @return connection of the id.
       */
      public Connection getConnection(String connectionId) {
          return connections.get(connectionId);
      }
      
      /**
       * get by client ip.
       *
       * @param clientIp client ip.
       * @return connections of the client ip.
       */
      public List<Connection> getConnectionByIp(String clientIp) {
          Set<Map.Entry<String, Connection>> entries = connections.entrySet();
          List<Connection> connections = new ArrayList<>();
          for (Map.Entry<String, Connection> entry : entries) {
              Connection value = entry.getValue();
              if (clientIp.equals(value.getMetaInfo().clientIp)) {
                  connections.add(value);
              }
          }
          return connections;
      }
      
      /**
       * init connection ejector.
       */
      public void initConnectionEjector() {
          String connectionRuntimeEjector = null;
          try {
              connectionRuntimeEjector = ControlConfigs.getInstance().getConnectionRuntimeEjector();
              Collection<RuntimeConnectionEjector> ejectors = NacosServiceLoader.load(RuntimeConnectionEjector.class);
              for (RuntimeConnectionEjector runtimeConnectionEjectorLoad : ejectors) {
                  if (runtimeConnectionEjectorLoad.getName().equalsIgnoreCase(connectionRuntimeEjector)) {
                      Loggers.CONNECTION.info("Found connection runtime ejector for name {}", connectionRuntimeEjector);
                      runtimeConnectionEjectorLoad.setConnectionManager(this);
                      runtimeConnectionEjector = runtimeConnectionEjectorLoad;
                  }
              }
          } catch (Throwable throwable) {
              Loggers.CONNECTION.warn("Fail to load  runtime ejector ", throwable);
          }
          
          if (runtimeConnectionEjector == null) {
              Loggers.CONNECTION
                      .info("Fail to find connection runtime ejector for name {},use default", connectionRuntimeEjector);
              NacosRuntimeConnectionEjector nacosRuntimeConnectionEjector = new NacosRuntimeConnectionEjector();
              nacosRuntimeConnectionEjector.setConnectionManager(this);
              runtimeConnectionEjector = nacosRuntimeConnectionEjector;
          }
      }
      
      /**
       * get current connections count.
       *
       * @return get all connection count
       */
      public int getCurrentConnectionCount() {
          return this.connections.size();
      }
      
      /**
       * refresh connection active time.
       *
       * @param connectionId connectionId.
       */
      public void refreshActiveTime(String connectionId) {
          Connection connection = connections.get(connectionId);
          if (connection != null) {
              connection.freshActiveTime();
          }
      }
      
      /**
       * Start Task：Expel the connection which active Time expire.
       */
      @PostConstruct
      public void start() {
          
          initConnectionEjector();
          // Start UnHealthy Connection Expel Task.
          RpcScheduledExecutor.COMMON_SERVER_EXECUTOR.scheduleWithFixedDelay(() -> {
              runtimeConnectionEjector.doEject();
              MetricsMonitor.getLongConnectionMonitor().set(connections.size());
          }, 1000L, 3000L, TimeUnit.MILLISECONDS);
  
          Boolean enabled = EnvUtil.getProperty("nacos.metric.grpc.server.connection.enabled", Boolean.class, true);
          if (enabled) {
              RpcScheduledExecutor.COMMON_SERVER_EXECUTOR.scheduleWithFixedDelay(() -> {
                  Map<String, Integer> count = new HashMap<>(16);
                  connections.forEach((id, connection) -> {
                      String module = connection.getLabels().getOrDefault(RemoteConstants.LABEL_MODULE, "unknown");
                      count.put(module, count.getOrDefault(module, 0) + 1);
                  });
                  MetricsMonitor.refreshModuleConnectionCount(count);
              }, 1L, EnvUtil.getProperty("nacos.metric.grpc.server.connection.interval", Long.class, 15L), TimeUnit.SECONDS);
          }
      }
      
      public void loadCount(int loadClient, String redirectAddress) {
          runtimeConnectionEjector.setLoadClient(loadClient);
          runtimeConnectionEjector.setRedirectAddress(redirectAddress);
      }
      
      /**
       * send load request to specific connectionId.
       *
       * @param connectionId    connection id of client.
       * @param redirectAddress server address to redirect.
       * @return whether remove connection.
       */
      public boolean loadSingle(String connectionId, String redirectAddress) {
          Connection connection = getConnection(connectionId);
          
          if (connection != null) {
              if (connection.getMetaInfo().isSdkSource()) {
                  ConnectResetRequest connectResetRequest = new ConnectResetRequest();
                  if (StringUtils.isNotBlank(redirectAddress) && redirectAddress.contains(Constants.COLON)) {
                      String[] split = redirectAddress.split(Constants.COLON);
                      connectResetRequest.setServerIp(split[0]);
                      connectResetRequest.setServerPort(split[1]);
                      connectResetRequest.setConnectionId(connectionId);
                  }
                  try {
                      connection.request(connectResetRequest, 3000L);
                  } catch (ConnectionAlreadyClosedException e) {
                      unregister(connectionId);
                  } catch (Exception e) {
                      LOGGER.error("error occurs when expel connection, connectionId: {} ", connectionId, e);
                      return false;
                  }
              }
          }
          return true;
          
      }
      
      /**
       * get all client count.
       *
       * @return client count.
       */
      public int currentClientsCount() {
          return connections.size();
      }
      
      /**
       * get client count with labels filter.
       *
       * @param filterLabels label to filter client count.
       * @return count with the specific filter labels.
       */
      public int currentClientsCount(Map<String, String> filterLabels) {
          int count = 0;
          for (Connection connection : connections.values()) {
              Map<String, String> labels = connection.getMetaInfo().labels;
              boolean disMatchFound = false;
              for (Map.Entry<String, String> entry : filterLabels.entrySet()) {
                  if (!entry.getValue().equals(labels.get(entry.getKey()))) {
                      disMatchFound = true;
                      break;
                  }
              }
              if (!disMatchFound) {
                  count++;
              }
          }
          return count;
      }
      
      /**
       * get client count from sdk.
       *
       * @return sdk client count.
       */
      public int currentSdkClientCount() {
          Map<String, String> filter = new HashMap<>(2);
          filter.put(RemoteConstants.LABEL_SOURCE, RemoteConstants.LABEL_SOURCE_SDK);
          return currentClientsCount(filter);
      }
      
      public Map<String, Connection> currentClients() {
          return connections;
      }
      
      public Map<String, AtomicInteger> getConnectionForClientIp() {
          return connectionForClientIp;
      }
  }
  