@Override
  public void put(SchedulerQueryContext query)
      throws OutOfCapacityException {
    Preconditions.checkNotNull(query);
    _queueLock.lock();
    String groupName = _groupSelector.getSchedulerGroupName(query);
    try {
      SchedulerGroup groupContext = getOrCreateGroupContext(groupName);
      checkGroupHasCapacity(groupContext);
      query.setSchedulerGroupContext(groupContext);
      groupContext.addLast(query);
      _queryReaderCondition.signal();
    } finally {
      _queueLock.unlock();
    }
  }
@Test
  public void testTakeWithLimits()
      throws OutOfCapacityException, BrokenBarrierException, InterruptedException {
    // Test that take() will not return query if that group is already using hardLimit resources
    Map<String, Object> properties = new HashMap<>();
    properties.put(ResourceManager.QUERY_WORKER_CONFIG_KEY, 40);
    properties.put(ResourceManager.QUERY_RUNNER_CONFIG_KEY, 10);
    properties.put(ResourceLimitPolicy.TABLE_THREADS_SOFT_LIMIT, 20);
    properties.put(ResourceLimitPolicy.TABLE_THREADS_HARD_LIMIT, 80);

    PinotConfiguration configuration = new PinotConfiguration(properties);

    PolicyBasedResourceManager rm = new PolicyBasedResourceManager(configuration);
    MultiLevelPriorityQueue queue = createQueue(configuration, rm);

    queue.put(createQueryRequest(GROUP_ONE, METRICS));
    queue.put(createQueryRequest(GROUP_ONE, METRICS));
    queue.put(createQueryRequest(GROUP_TWO, METRICS));
    // group one has higher priority but it's above soft thread limit
    TestSchedulerGroup testGroupOne = GROUP_FACTORY._groupMap.get(GROUP_ONE);
    TestSchedulerGroup testGroupTwo = GROUP_FACTORY._groupMap.get(GROUP_TWO);

    testGroupOne.addReservedThreads(rm.getTableThreadsSoftLimit() + 1);
    QueueReader reader = new QueueReader(queue);
    reader.startAndWaitForRead();
    assertEquals(reader._readQueries.size(), 1);
    assertEquals(reader._readQueries.poll().getSchedulerGroup().name(), GROUP_TWO);

    // add one more group two
    queue.put(createQueryRequest(GROUP_TWO, METRICS));
    reader = new QueueReader(queue);
    reader.startAndWaitForRead();
    assertEquals(reader._readQueries.size(), 1);
    assertEquals(reader._readQueries.poll().getSchedulerGroup().name(), GROUP_TWO);

    // add one more groupTwo and set groupTwo threads to higher than groupOne
    queue.put(createQueryRequest(GROUP_TWO, METRICS));
    testGroupTwo.addReservedThreads(testGroupOne.totalReservedThreads() + 1);
    reader = new QueueReader(queue);
    reader.startAndWaitForRead();
    assertEquals(reader._readQueries.size(), 1);
    assertEquals(reader._readQueries.poll().getSchedulerGroup().name(), GROUP_ONE);

    // set groupOne above hard limit
    testGroupOne.addReservedThreads(rm.getTableThreadsHardLimit());
    reader = new QueueReader(queue);
    reader.startAndWaitForRead();
    assertEquals(reader._readQueries.size(), 1);
    assertEquals(reader._readQueries.poll().getSchedulerGroup().name(), GROUP_TWO);

    // all groups above hard limit
    queue.put(createQueryRequest(GROUP_TWO, METRICS));
    queue.put(createQueryRequest(GROUP_TWO, METRICS));
    queue.put(createQueryRequest(GROUP_ONE, METRICS));
    testGroupTwo.addReservedThreads(rm.getTableThreadsHardLimit());
    reader = new QueueReader(queue);
    reader.startAndWaitForQueueWakeup();
    assertEquals(reader._readQueries.size(), 0);
    // try again
    sleepForQueueWakeup(queue);
    assertEquals(reader._readQueries.size(), 0);

    // now set thread limit lower for a group (aka. query finished)
    testGroupTwo.releasedReservedThreads(testGroupTwo.totalReservedThreads());
    sleepForQueueWakeup(queue);
    assertEquals(reader._readQueries.size(), 1);
  }