content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You 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 org.apache.kafka.coordinator.common.runtime;
  
  import org.apache.kafka.common.utils.Time;
  
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.PriorityQueue;
  import java.util.concurrent.TimeUnit;
  
  /**
   * A simple mock for the {@link CoordinatorTimer}. The mock does not automatically
   * expire timeouts. They are only expired when {@link MockCoordinatorTimer#poll()}
   * is called.
   */
  public class MockCoordinatorTimer<T, U> implements CoordinatorTimer<T, U> {
      /**
       * Represents a scheduled timeout.
       */
      public static class ScheduledTimeout<T, U> {
          public final String key;
          public final long deadlineMs;
          public final TimeoutOperation<T, U> operation;
  
          public ScheduledTimeout(
              String key,
              long deadlineMs,
              TimeoutOperation<T, U> operation
          ) {
              this.key = key;
              this.deadlineMs = deadlineMs;
              this.operation = operation;
          }
      }
  
      /**
       * Represents an expired timeout.
       */
      public static class ExpiredTimeout<T, U> {
          public final String key;
          public final CoordinatorResult<T, U> result;
  
          public ExpiredTimeout(
              String key,
              CoordinatorResult<T, U> result
          ) {
              this.key = key;
              this.result = result;
          }
  
          @Override
          public boolean equals(Object o) {
              if (this == o) return true;
              if (o == null || getClass() != o.getClass()) return false;
  
              ExpiredTimeout<?, ?> that = (ExpiredTimeout<?, ?>) o;
  
              if (!Objects.equals(key, that.key)) return false;
              return Objects.equals(result, that.result);
          }
  
          @Override
          public int hashCode() {
              int result1 = key != null ? key.hashCode() : 0;
              result1 = 31 * result1 + (result != null ? result.hashCode() : 0);
              return result1;
          }
      }
  
      private final Time time;
  
      private final Map<String, ScheduledTimeout<T, U>> timeoutMap = new HashMap<>();
      private final PriorityQueue<ScheduledTimeout<T, U>> timeoutQueue = new PriorityQueue<>(
          Comparator.comparingLong(entry -> entry.deadlineMs)
      );
  
      public MockCoordinatorTimer(Time time) {
          this.time = time;
      }
  
      /**
       * Schedules a timeout.
       */
      @Override
      public void schedule(
          String key,
          long delay,
          TimeUnit unit,
          boolean retry,
          long retryBackoff,
          TimeoutOperation<T, U> operation
      ) {
          cancel(key);
  
          long deadlineMs = time.milliseconds() + unit.toMillis(delay);
          ScheduledTimeout<T, U> timeout = new ScheduledTimeout<>(key, deadlineMs, operation);
          timeoutQueue.add(timeout);
          timeoutMap.put(key, timeout);
      }
  
      @Override
      public void schedule(
          String key,
          long delay,
          TimeUnit unit,
          boolean retry,
          TimeoutOperation<T, U> operation
      ) {
          schedule(key, delay, unit, retry, 500L, operation);
      }
  
      @Override
      public void scheduleIfAbsent(
          String key,
          long delay,
          TimeUnit unit,
          boolean retry,
          TimeoutOperation<T, U> operation
      ) {
          if (!timeoutMap.containsKey(key)) {
              schedule(key, delay, unit, retry, 500L, operation);
          }
      }
  
      /**
       * Cancels a timeout.
       */
      @Override
      public void cancel(String key) {
          ScheduledTimeout<T, U> timeout = timeoutMap.remove(key);
          if (timeout != null) {
              timeoutQueue.remove(timeout);
          }
      }
  
      /**
       * @return True if a timeout with the key exists; false otherwise.
       */
      public boolean contains(String key) {
          return timeoutMap.containsKey(key);
      }
  
      /**
       * @return The scheduled timeout for the key; null otherwise.
       */
      public ScheduledTimeout<T, U> timeout(String key) {
          return timeoutMap.get(key);
      }
  
      /**
       * @return The number of scheduled timeouts.
       */
      public int size() {
          return timeoutMap.size();
      }
  
      /**
       * @return A list of expired timeouts based on the current time.
       */
      public List<ExpiredTimeout<T, U>> poll() {
          List<ExpiredTimeout<T, U>> results = new ArrayList<>();
  
          ScheduledTimeout<T, U> timeout = timeoutQueue.peek();
          while (timeout != null && timeout.deadlineMs <= time.milliseconds()) {
              timeoutQueue.poll();
              timeoutMap.remove(timeout.key, timeout);
  
              results.add(new ExpiredTimeout<>(
                  timeout.key,
                  timeout.operation.generateRecords()
              ));
  
              timeout = timeoutQueue.peek();
          }
  
          return results;
      }
  }
  