package org.geoserver.flow.controller;
import java.rmi.server.UID;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.servlet.http.Cookie;
import org.geoserver.flow.ControlFlowCallback;
import org.geoserver.flow.FlowController;
import org.geoserver.ows.Request;
import org.geotools.util.logging.Logging;
/**
* A flow controller setting a cookie on HTTP request and making sure the same user cannot do more
* than X requests in parallel. Warning: if a client does not support cookies this class cannot work
* properly and will start accumulating queues with just one item inside. As a workaround when too
* many queues are accumulated a scan starts that purges all queues that are empty and have not been
* touched within a given amount of time: the idea is that a past that time we're assuming the
* client is no more working actively against the server and the queue can thus be removed.
*
* @author Andrea Aime - OpenGeo
*
*/
public class UserFlowController implements FlowController {
static final Logger LOGGER = Logging.getLogger(ControlFlowCallback.class);
static String COOKIE_NAME = "GS_FLOW_CONTROL";
static String COOKIE_PREFIX = "GS_CFLOW_";
/**
* Thread local holding the current request queue id TODO: consider having a user map in
* {@link Request} instead
*/
static ThreadLocal<String> QUEUE_ID = new ThreadLocal<String>();
/**
* The size of each queue
*/
int queueSize;
/**
* Last time we've performed a queue cleanup
*/
long lastCleanup = System.currentTimeMillis();
/**
* Number of queues at which we start looking for purging stale ones
*/
int maxQueues = 100;
/**
* Time it takes for an inactive queue to be considered stale
*/
int maxAge = 10000;
/**
* The per user queue collection
*/
Map<String, TimedBlockingQueue> queues = new ConcurrentHashMap<String, TimedBlockingQueue>();
/**
* Builds a UserFlowController that will trigger stale queue expiration once 100 queues have
* been accumulated and
*
* @param queueSize
* the maximum amount of per user concurrent requests
*/
public UserFlowController(int queueSize) {
this(queueSize, 100, 10000);
}
/**
* Builds a new {@link UserFlowController}
*
* @param queueSize
* the maximum amount of per user concurrent requests
* @param maxQueues
* the number of accumulated user queues that will trigger a queue cleanup
* @param maxAge
* the max quiet time for an empty queue to be considered stale and removed
*/
public UserFlowController(int queueSize, int maxQueues, int maxAge) {
this.queueSize = queueSize;
this.maxQueues = maxQueues;
this.maxAge = maxAge;
}
public int getPriority() {
return queueSize;
}
public void requestComplete(Request request) {
String queueId = QUEUE_ID.get();
QUEUE_ID.remove();
BlockingQueue<Request> queue = queues.get(queueId);
if (queue != null)
queue.remove(request);
}
public boolean requestIncoming(Request request, long timeout) {
boolean retval = true;
long now = System.currentTimeMillis();
// check if this client already made other connections
Cookie idCookie = null;
Cookie[] cookies = request.getHttpRequest().getCookies();
if (cookies != null) {
for (Cookie cookie : cookies) {
if (cookie.getName().equals(COOKIE_NAME)) {
idCookie = cookie;
break;
}
}
}
// see if we have that queue already
TimedBlockingQueue queue = null;
if (idCookie != null) {
queue = queues.get(idCookie.getValue());
}
// generate a unique queue id for this client if none was found
if(queue == null) {
idCookie = new Cookie(COOKIE_NAME, COOKIE_PREFIX + new UID().toString());
queue = new TimedBlockingQueue(queueSize, true);
queues.put(idCookie.getValue(), queue);
}
QUEUE_ID.set(idCookie.getValue());
request.getHttpResponse().addCookie(idCookie);
// queue token handling
try {
if(timeout > 0) {
retval = queue.offer(request, timeout, TimeUnit.MILLISECONDS);
} else {
queue.put(request);
}
} catch (InterruptedException e) {
LOGGER.log(Level.WARNING, "Unexpected interruption while "
+ "blocking on the request queue");
}
if (LOGGER.isLoggable(Level.FINE)) {
LOGGER.fine("UserFlowController(" + queueSize + "," + idCookie.getValue()
+ ") queue size " + queue.size());
LOGGER.fine("UserFlowController(" + queueSize + "," + idCookie.getValue()
+ ") total queues " + queues.size());
}
// cleanup stale queues if necessary
if ((queues.size() > maxQueues && (now - lastCleanup) > (maxAge / 10))
|| (now - lastCleanup) > maxAge) {
int cleanupCount = 0;
synchronized (queues) {
for (String key : queues.keySet()) {
TimedBlockingQueue tbq = queues.get(key);
if (now - tbq.lastModified > maxAge && tbq.size() == 0) {
queues.remove(key);
cleanupCount++;
}
}
lastCleanup = now;
if (LOGGER.isLoggable(Level.FINE)) {
LOGGER.fine("UserFlowController(" + queueSize + ") purged " + cleanupCount
+ " stale queues");
}
}
}
return retval;
}
@SuppressWarnings("serial")
static class TimedBlockingQueue extends ArrayBlockingQueue<Request> {
long lastModified;
public TimedBlockingQueue(int capacity, boolean fair) {
super(capacity, fair);
}
@Override
public void put(Request o) throws InterruptedException {
super.put(o);
lastModified = System.currentTimeMillis();
}
@Override
public boolean remove(Object o) {
lastModified = System.currentTimeMillis();
return super.remove(o);
}
}
}