How to implement fair usage policy of threads in API backend?

Question

I am building an REST API in Java which I would be exposing to the outside world. People who would be invoking the API would have to be registered and would be sending their userid in the request. There would be a maximum of, say, 10 concurrent threads available for executing the API request. I am maintaining a queue which holds all the request ids to be serviced (the primary key of the DB entry). I need to implement some fair usage policy as follows - If there are more than 10 jobs in the queue (i.e more than max number of threads), a user is allowed to execute only one request at a time (the other requests submitted by him/her, if any, would remain in the queue and would be taken up only once his previous request has completed execution). In case there are free threads, i.e. even after allotting threads to requests submitted by different users, then the remaining threads in the thread pool can be distributed among the remaining requests (even if the user who has submitted the request is already holding one thread at that moment).

The current implementation is as follows -

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.PriorityBlockingQueue;
import java.util.concurrent.Semaphore;


public class APIJobExecutor implements Runnable{
    private static PriorityBlockingQueue<Integer> jobQueue = new PriorityBlockingQueue<Integer>();
    private static ExecutorService jobExecutor = Executors.newCachedThreadPool();
    private static final int MAX_THREADS = 10;
    private static Semaphore sem = new Semaphore(MAX_THREADS, true);

    private APIJobExecutor(){

    }

    public static void addJob(int jobId)
    {
        if(!jobQueue.contains(jobId)){
            jobQueue.add(new Integer(jobId));
        }
    }

    public void run()
    {
        while (true) {
            try {
                sem.acquire();
            }catch (InterruptedException e1) {
                e1.printStackTrace();
                //unable to acquire lock. retry.
                continue;
            }
            try {
                Integer jobItem = jobQueue.take();
                jobExecutor.submit(new APIJobService(jobItem));
            } catch (InterruptedException e) {
                e.printStackTrace();
            } finally {
                sem.release();
            }
        }
    }
}

Edit: Is there any out of the box Java data structure that gives me this functionality. If not, how do I go about implementing the same?

Answer 1

This is a fairly common "quality of service" pattern and can be solved using the bucket idea within a job-queue. I do not know of a standard Java implementation and/or datastructure for this pattern (maybe the PriorityQueue?), but there should be at least a couple of implementations available (let us know if you find a good one).

I did once create my own implementation and I've tried to de-couple it from the project so that you may modify and use it (add unit-tests!). A couple of notes:

• a default-queue is used in case QoS is not needed (e.g. if less than 10 jobs are executing).
• the basic idea is to store tasks in lists per QoS-key (e.g. the username), and maintain a separate "who is next" list.
• it is intended to be used within a job queue (e.g. part of the APIJobExecutor, not a replacement). Part of the job queue's responsibility is to always call remove(taskId) after a task is executed.
• there should be no memory leaks: if there are no tasks/jobs in the queue, all internal maps and lists should be empty.

The code:

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;
import java.util.*;

import org.slf4j.*;

/** A FIFO task queue. */
public class QosTaskQueue<TASKTYPE, TASKIDTYPE> {
private static final Logger log = LoggerFactory.getLogger(QosTaskQueue.class);

public static final String EMPTY_STRING = "";

/** Work tasks queued which have no (relevant) QoS key. */
private final ConcurrentLinkedQueue<TASKIDTYPE> defaultQ = new ConcurrentLinkedQueue<TASKIDTYPE>();

private final AtomicInteger taskQSize = new AtomicInteger();

private final Map<TASKIDTYPE, TASKTYPE> queuedTasks = new ConcurrentHashMap<TASKIDTYPE, TASKTYPE>();

/** Amount of tasks in queue before "quality of service" distribution kicks in. */
private int qosThreshold = 10;

/** Indicates if "quality of service" distribution is in effect. */
private volatile boolean usingQos;

/**
 * Lock for all modifications to Qos-queues.
 * <br>Must be "fair" to ensure adding does not block polling threads forever and vice versa. 
 */
private final ReentrantLock qosKeyLock = new ReentrantLock(true); 

/*
 * Since all QoS modifications can be done by multiple threads simultaneously, 
 * there is never a good time to add or remove a Qos-key with associated queue. 
 * There is always a chance that a key is added while being removed and vice versa.
 * The simplest solution is to make everything synchronized, which is what qosKeyLock is used for.
 */
private final Map<String, Queue<TASKIDTYPE>> qosQueues = new HashMap<String, Queue<TASKIDTYPE>>();
private final Queue<String> qosTurn = new LinkedList<String>();

public boolean add(TASKTYPE wt, TASKIDTYPE taskId, String qosKey) {

    if (queuedTasks.containsKey(taskId))  {
        throw new IllegalStateException("Task with ID [" + taskId + "] already enqueued.");
    }
    queuedTasks.put(taskId, wt);
    return addToQ(taskId, qosKey);
}

public TASKTYPE poll() {

    TASKIDTYPE taskId = pollQos();
    return (taskId == null ? null : queuedTasks.get(taskId));
}

/**
 * This method must be called after a task is taken from the queue
 * using {@link #poll()} and executed.
 */
public TASKTYPE remove(TASKIDTYPE taskId) {

    TASKTYPE wt = queuedTasks.remove(taskId);
    if (wt != null) {
        taskQSize.decrementAndGet();
    }
    return wt;
}

private boolean addToQ(TASKIDTYPE taskId, String qosKey) {

    if (qosKey == null || qosKey.equals(EMPTY_STRING) || size() < getQosThreshold()) {
        defaultQ.add(taskId);
    } else {
        addSynced(taskId, qosKey);
    }
    taskQSize.incrementAndGet();
    return true;
}

private void addSynced(TASKIDTYPE taskId, String qosKey) {

    qosKeyLock.lock();
    try {
        Queue<TASKIDTYPE> qosQ = qosQueues.get(qosKey); 
        if (qosQ == null) {
            if (!isUsingQos()) {
                // Setup QoS mechanics
                qosTurn.clear();
                qosTurn.add(EMPTY_STRING);
                usingQos = true;
            }
            qosQ = new LinkedList<TASKIDTYPE>();
            qosQ.add(taskId);
            qosQueues.put(qosKey, qosQ);
            qosTurn.add(qosKey);
            log.trace("Created QoS queue for {}", qosKey);
        } else { 
            qosQ.add(taskId);
            if (log.isTraceEnabled()) {
                log.trace("Added task to QoS queue {}, size: " + qosQ.size(), qosKey);
            }
        }
    } finally {
        qosKeyLock.unlock();
    }
}

private TASKIDTYPE pollQos() {

    TASKIDTYPE taskId = null;
    qosKeyLock.lock();
    try {
        taskId = pollQosRecursive();
    } finally {
        qosKeyLock.unlock();
    }
    return taskId;
}

/**
 * Poll the work task queues according to qosTurn.
 * Recursive in case empty QoS queues are removed or defaultQ is empty.  
 * @return
 */
private TASKIDTYPE pollQosRecursive() {

    if (!isUsingQos()) {
        // QoS might have been disabled before lock was released or by this recursive method.
        return defaultQ.poll();
    }
    String qosKey = qosTurn.poll();
    Queue<TASKIDTYPE> qosQ = (qosKey.equals(EMPTY_STRING) ? defaultQ : qosQueues.get(qosKey));
    TASKIDTYPE taskId = qosQ.poll();
    if (qosQ == defaultQ) {
        // DefaultQ should always be checked, even if it was empty
        qosTurn.add(EMPTY_STRING);
        if (taskId == null) {
            taskId = pollQosRecursive();
        } else {
            log.trace("Removed task from defaultQ.");
        }
    } else {
        if (taskId == null) {
            qosQueues.remove(qosKey);
            if (qosQueues.isEmpty()) {
                usingQos = false;
            }
            taskId = pollQosRecursive();
        } else {
            qosTurn.add(qosKey);
            if (log.isTraceEnabled()) {
                log.trace("Removed task from QoS queue {}, size: " + qosQ.size(), qosKey);
            }
        }
    }
    return taskId;
}

@Override
public String toString() {

    StringBuilder sb = new StringBuilder(this.getClass().getName());
    sb.append(", size: ").append(size());
    sb.append(", number of QoS queues: ").append(qosQueues.size());
    return sb.toString();
}

public boolean containsTaskId(TASKIDTYPE wid) {
    return queuedTasks.containsKey(wid);
}

public int size() { 
    return taskQSize.get(); 
}

public void setQosThreshold(int size) {
    this.qosThreshold = size;
}

public int getQosThreshold() {
    return qosThreshold;
}

public boolean isUsingQos() {
    return usingQos;
}

}