/**
 * Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com)
 *
 * 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.linkedin.pinot.transport.common;

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

import org.apache.http.annotation.ThreadSafe;

import com.linkedin.pinot.common.response.ServerInstance;


/**
 * Maintains next pointer per segment basis. Expected to be thread-safe
 *
 */
@ThreadSafe
public class RoundRobinReplicaSelection extends ReplicaSelection {

  private final Map<SegmentId, AtomicInteger> _nextPositionMap;

  public RoundRobinReplicaSelection() {
    _nextPositionMap = new ConcurrentHashMap<SegmentId, AtomicInteger>();
  }

  @Override
  public void reset(SegmentId p) {
    _nextPositionMap.remove(p);
  }

  @Override
  public void reset(SegmentIdSet pg) {
    for (SegmentId p : pg.getSegments()) {
      reset(p);
    }
  }

  @Override
  public ServerInstance selectServer(SegmentId p, List<ServerInstance> orderedServers, Object bucketKey) {

    int size = orderedServers.size();

    if (size <= 0) {
      return null;
    }

    AtomicInteger a = _nextPositionMap.get(p);

    // Create a new position tracker for the partition if it is not available
    if (null == a) {
      synchronized (this) {
        a = _nextPositionMap.get(p);
        if (null == a) {
          a = new AtomicInteger(-1);
          _nextPositionMap.put(p, a);
        }
      }
    }

    return orderedServers.get(Math.abs(a.incrementAndGet()) % size);
  }

}