com.rabbitmq.client.impl.WorkPool Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of amqp-client Show documentation
Show all versions of amqp-client Show documentation
The RabbitMQ Java client library allows Java applications to interface with RabbitMQ.
// Copyright (c) 2007-2023 Broadcom. All Rights Reserved. The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
//
// This software, the RabbitMQ Java client library, is triple-licensed under the
// Mozilla Public License 2.0 ("MPL"), the GNU General Public License version 2
// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2. For the ASL,
// please see LICENSE-APACHE2.
//
// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
// either express or implied. See the LICENSE file for specific language governing
// rights and limitations of this software.
//
// If you have any questions regarding licensing, please contact us at
// [email protected].
package com.rabbitmq.client.impl;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
/**
* This is a generic implementation of the channels specification
* in Channeling Work, Nov 2010 (channels.pdf).
* Objects of type K must be registered, with registerKey(K),
* and then they become clients and a queue of
* items (type W) is stored for each client.
*
* Each client has a state which is exactly one of dormant,
* in progress or ready. Immediately after registration a client is dormant.
* Items may be (singly) added to (the end of) a client's queue with {@link WorkPool#addWorkItem(Object, Object)}.
* If the client is dormant it becomes ready thereby. All other states remain unchanged.
* The next ready client, together with a collection of its items,
* may be retrieved with nextWorkBlock(collection,max)
* (making that client in progress).
* An in progress client can finish (processing a batch of items) with finishWorkBlock(K).
* It then becomes either dormant or ready, depending if its queue of work items is empty or no.
* If a client has items queued, it is either in progress or ready but cannot be both.
* When work is finished it may be marked ready if there is further work,
* or dormant if there is not.
* There is never any work for a dormant client.
* A client may be unregistered, with unregisterKey(K), which removes the client from
* all parts of the state, and any queue of items stored with it.
* All clients may be unregistered with unregisterAllKeys().
* Concurrent Semantics
* This implementation is thread-safe.
* @param Key -- type of client
* @param Work -- type of work item
*/
public class WorkPool {
private static final int MAX_QUEUE_LENGTH = 1000;
/** An injective queue of ready clients. */
private final SetQueue ready = new SetQueue();
/** The set of clients which have work in progress. */
private final Set inProgress = new HashSet();
/** The pool of registered clients, with their work queues. */
private final Map> pool = new HashMap>();
/** Those keys which want limits to be removed. We do not limit queue size if this is non-empty. */
private final Set unlimited = new HashSet();
private final BiConsumer, W> enqueueingCallback;
public WorkPool(final int queueingTimeout) {
if (queueingTimeout > 0) {
this.enqueueingCallback = (queue, item) -> {
try {
boolean offered = queue.offer(item, queueingTimeout, TimeUnit.MILLISECONDS);
if (!offered) {
throw new WorkPoolFullException("Could not enqueue in work pool after " + queueingTimeout + " ms.");
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
};
} else {
this.enqueueingCallback = (queue, item) -> {
try {
queue.put(item);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
};
}
}
/**
* Add client key to pool of item queues, with an empty queue.
* A client is initially dormant.
* No-op if key already present.
* @param key client to add to pool
*/
public void registerKey(K key) {
synchronized (this) {
if (!this.pool.containsKey(key)) {
int initialCapacity = unlimited.isEmpty() ? MAX_QUEUE_LENGTH : Integer.MAX_VALUE;
this.pool.put(key, new VariableLinkedBlockingQueue(initialCapacity));
}
}
}
public synchronized void limit(K key) {
unlimited.remove(key);
if (unlimited.isEmpty()) {
setCapacities(MAX_QUEUE_LENGTH);
}
}
public synchronized void unlimit(K key) {
unlimited.add(key);
if (!unlimited.isEmpty()) {
setCapacities(Integer.MAX_VALUE);
}
}
private void setCapacities(int capacity) {
Iterator> it = pool.values().iterator();
while (it.hasNext()) {
it.next().setCapacity(capacity);
}
}
/**
* Remove client from pool and from any other state. Has no effect if client already absent.
* @param key of client to unregister
*/
public void unregisterKey(K key) {
synchronized (this) {
this.pool.remove(key);
this.ready.remove(key);
this.inProgress.remove(key);
this.unlimited.remove(key);
}
}
/**
* Remove all clients from pool and from any other state.
*/
public void unregisterAllKeys() {
synchronized (this) {
this.pool.clear();
this.ready.clear();
this.inProgress.clear();
this.unlimited.clear();
}
}
/**
* Return the next ready client,
* and transfer a collection of that client's items to process.
* Mark client in progress.
* If there is no ready client, return null.
* @param to collection object in which to transfer items
* @param size max number of items to transfer
* @return key of client to whom items belong, or null if there is none.
*/
public K nextWorkBlock(Collection to, int size) {
synchronized (this) {
K nextKey = readyToInProgress();
if (nextKey != null) {
VariableLinkedBlockingQueue queue = this.pool.get(nextKey);
drainTo(queue, to, size);
}
return nextKey;
}
}
/**
* Private implementation of drainTo (not implemented for LinkedList<W>s).
* @param deList to take (poll) elements from
* @param c to add elements to
* @param maxElements to take from deList
* @return number of elements actually taken
*/
private int drainTo(VariableLinkedBlockingQueue deList, Collection c, int maxElements) {
int n = 0;
while (n < maxElements) {
W first = deList.poll();
if (first == null)
break;
c.add(first);
++n;
}
return n;
}
/**
* Add (enqueue) an item for a specific client.
* No change and returns false if client not registered.
* If dormant, the client will be marked ready.
* @param key the client to add to the work item to
* @param item the work item to add to the client queue
* @return true if and only if the client is marked ready
* — as a result of this work item
*/
public boolean addWorkItem(K key, W item) {
VariableLinkedBlockingQueue queue;
synchronized (this) {
queue = this.pool.get(key);
}
// The put operation may block. We need to make sure we are not holding the lock while that happens.
if (queue != null) {
enqueueingCallback.accept(queue, item);
synchronized (this) {
if (isDormant(key)) {
dormantToReady(key);
return true;
}
}
}
return false;
}
/**
* Set client no longer in progress.
* Ignore unknown clients (and return false).
* @param key client that has finished work
* @return true if and only if client becomes ready
* @throws IllegalStateException if registered client not in progress
*/
public boolean finishWorkBlock(K key) {
synchronized (this) {
if (!this.isRegistered(key))
return false;
if (!this.inProgress.contains(key)) {
throw new IllegalStateException("Client " + key + " not in progress");
}
if (moreWorkItems(key)) {
inProgressToReady(key);
return true;
} else {
inProgressToDormant(key);
return false;
}
}
}
private boolean moreWorkItems(K key) {
VariableLinkedBlockingQueue leList = this.pool.get(key);
return leList != null && !leList.isEmpty();
}
/* State identification functions */
private boolean isInProgress(K key){ return this.inProgress.contains(key); }
private boolean isReady(K key){ return this.ready.contains(key); }
private boolean isRegistered(K key) { return this.pool.containsKey(key); }
private boolean isDormant(K key){ return !isInProgress(key) && !isReady(key) && isRegistered(key); }
/* State transition methods - all assume key registered */
private void inProgressToReady(K key){ this.inProgress.remove(key); this.ready.addIfNotPresent(key); }
private void inProgressToDormant(K key){ this.inProgress.remove(key); }
private void dormantToReady(K key){ this.ready.addIfNotPresent(key); }
/* Basic work selector and state transition step */
private K readyToInProgress() {
K key = this.ready.poll();
if (key != null) {
this.inProgress.add(key);
}
return key;
}
}