/**
* 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
* <p/>
* http://www.apache.org/licenses/LICENSE-2.0
* <p/>
* 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 backtype.storm.utils;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import com.alibaba.jstorm.callback.Callback;
import com.alibaba.jstorm.daemon.worker.Flusher;
import com.alibaba.jstorm.utils.JStormUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.lmax.disruptor.AlertException;
import com.lmax.disruptor.EventFactory;
import com.lmax.disruptor.EventHandler;
import com.lmax.disruptor.InsufficientCapacityException;
import com.lmax.disruptor.RingBuffer;
import com.lmax.disruptor.Sequence;
import com.lmax.disruptor.SequenceBarrier;
import com.lmax.disruptor.TimeoutException;
import com.lmax.disruptor.WaitStrategy;
import com.lmax.disruptor.dsl.ProducerType;
/**
*
* A single consumer queue that uses the LMAX Disruptor. They key to the performance is the ability to catch up to the producer by processing tuples in batches.
*/
public class DisruptorQueueImpl extends DisruptorQueue {
private static final Logger LOG = LoggerFactory.getLogger(DisruptorQueueImpl.class);
private static final Object INTERRUPT = new Object();
private static final String PREFIX = "disruptor-";
private final String _queueName;
private final RingBuffer<MutableObject> _buffer;
private final Sequence _consumer;
private final SequenceBarrier _barrier;
private boolean _isBatch;
private ThreadLocalBatch _batcher;
private int _inputBatchSize = 0;
private Flusher _flusher;
private Object _callbackLock = new Object();
private List<Callback> _callbacks = null;
private Object _cacheLock = new Object();
private List<Object> _cache;
// TODO: consider having a threadlocal cache of this variable to improve performance
// reads?
// volatile boolean consumerStartedFlag = false;
private final HashMap<String, Object> state = new HashMap<>(4);
public DisruptorQueueImpl(String queueName, ProducerType producerType, int bufferSize, WaitStrategy wait, boolean isBatch, int batchSize, long flushMs) {
_queueName = PREFIX + queueName;
_buffer = RingBuffer.create(producerType, new ObjectEventFactory(), bufferSize, wait);
_consumer = new Sequence();
_barrier = _buffer.newBarrier();
_buffer.addGatingSequences(_consumer);
_isBatch = isBatch;
_cache = new ArrayList<>();
_inputBatchSize = batchSize;
if (_isBatch) {
_batcher = new ThreadLocalBatch();
_flusher = new DisruptorFlusher(Math.max(flushMs, 1));
_flusher.start();
} else {
_batcher = null;
}
}
public String getName() {
return _queueName;
}
public void consumeBatch(EventHandler<Object> handler) {
// write pos > read pos
// Asynchronous release the queue, but still is single thread
if (_buffer.getCursor() > _consumer.get())
consumeBatchWhenAvailable(handler);
}
public void haltWithInterrupt() {
publish(INTERRUPT);
}
public Object poll() {
return getEvent(false);
}
public Object take() {
return getEvent(true);
}
private Object getEvent(boolean block) {
Object ret = null;
if (cacheSize() > 0) {
synchronized (_cacheLock) {
ret = _cache.remove(0);
}
} else {
if (block || _buffer.getCursor() > _consumer.get()) {
final long nextSequence = _consumer.get() + 1;
try {
final long availableSequence = _barrier.waitFor(nextSequence);
if (availableSequence >= nextSequence) {
MutableObject mo = _buffer.get(nextSequence);
_consumer.set(nextSequence);
ret = mo.o;
mo.setObject(null);
}
} catch (AlertException e) {
LOG.error(e.getMessage(), e);
throw new RuntimeException(e);
} catch (InterruptedException e) {
LOG.error("InterruptedException " + e.getCause());
// throw new RuntimeException(e);
return null;
} catch (TimeoutException e) {
// LOG.error(e.getCause(), e);
return null;
}
}
}
handlerCallback();
return ret;
}
private void handlerCallback() {
if (_callbacks != null) {
synchronized (_callbackLock) {
if (_callbacks != null) {
Iterator<Callback> itr = _callbacks.iterator();
while (itr.hasNext()) {
Callback cb = itr.next();
if ((boolean) cb.execute(this)) {
itr.remove();
}
}
if (_callbacks.size() == 0) {
_callbacks = null;
}
}
}
}
}
public void consumeBatchWhenAvailable(EventHandler<Object> handler) {
consumeBatchWhenAvailable(handler, true);
}
public void consumeBatchWhenAvailableWithCallback(EventHandler<Object> handler) {
consumeBatchWhenAvailable(handler);
handlerCallback();
}
public void multiConsumeBatchWhenAvailable(EventHandler<Object> handler) {
consumeBatchWhenAvailable(handler, false);
}
public void multiConsumeBatchWhenAvailableWithCallback(EventHandler<Object> handler) {
consumeBatchWhenAvailable(handler, false);
handlerCallback();
}
public void consumeBatchWhenAvailable(EventHandler<Object> handler, boolean isSync) {
List<Object> cache = null;
synchronized (_cacheLock) {
if (_cache.size() > 0) {
cache = _cache;
_cache = new ArrayList<>();
}
}
if (cache != null) {
for (int i = 0; i < cache.size(); i++) {
try {
handler.onEvent(cache.get(i), 0, i == (cache.size() - 1));
} catch (Exception e) {
LOG.error(e.getMessage(), e);
throw new RuntimeException(e);
}
}
} else {
try {
if (isSync)
consumeBatchToCursor(handler);
else
asyncConsumeBatchToCursor(handler);
} catch (AlertException e) {
LOG.error(e.getMessage(), e);
throw new RuntimeException(e);
} catch (InterruptedException e) {
LOG.error("InterruptedException " + e.getCause());
} catch (TimeoutException e) {
// Do nothing
}
}
}
private long getAvailableConsumeCursor() throws AlertException, InterruptedException, TimeoutException {
final long nextSequence = _consumer.get() + 1;
return _barrier.waitFor(nextSequence);
}
public void consumeBatchToCursor(EventHandler<Object> handler) throws AlertException, InterruptedException, TimeoutException {
long endCursor = getAvailableConsumeCursor();
long curr = _consumer.get() + 1;
for (; curr <= endCursor; curr++) {
try {
MutableObject mo = _buffer.get(curr);
Object o = mo.o;
mo.setObject(null);
handler.onEvent(o, curr, curr == endCursor);
} catch (InterruptedException e) {
LOG.error(e.getMessage(), e);
return;
} catch (Exception e) {
LOG.error(e.getMessage(), e);
throw new RuntimeException(e);
}
}
// TODO: only set this if the consumer cursor has changed?
_consumer.set(endCursor);
}
public void asyncConsumeBatchToCursor(EventHandler<Object> handler) throws AlertException, InterruptedException, TimeoutException {
List<Object> batch = getConsumeBatch();
if (batch == null)
return;
for (int i = 0; i < batch.size(); i++) {
try {
handler.onEvent(batch.get(i), 0, i == (batch.size() - 1));
} catch (Exception e) {
LOG.error(e.getMessage(), e);
throw new RuntimeException(e);
}
}
}
private synchronized List<Object> getConsumeBatch() throws AlertException, InterruptedException, TimeoutException {
long endCursor = getAvailableConsumeCursor();
long currCursor = _consumer.get();
long eventNumber = endCursor - currCursor;
List<Object> batch = new ArrayList<>((int) eventNumber);
for (long curr = currCursor + 1; curr <= endCursor; curr++) {
try {
MutableObject mo = _buffer.get(curr);
Object o = mo.o;
mo.setObject(null);
batch.add(o);
} catch (Exception e) {
LOG.error(e.getMessage(), e);
throw new RuntimeException(e);
}
}
_consumer.set(endCursor);
return batch;
}
@Override
public List<Object> retreiveAvailableBatch() throws AlertException, InterruptedException, TimeoutException {
// get all events in disruptor queue
return getConsumeBatch();
}
/*
* Caches until consumerStarted is called, upon which the cache is flushed to the consumer
*/
public void publish(Object obj) {
try {
if (_isBatch) {
publishBatch(obj);
} else {
publish(obj, true);
}
} catch (InsufficientCapacityException ex) {
throw new RuntimeException("This code should be unreachable!");
}
}
public void publishBatch(Object obj) {
_batcher.addAndFlush(obj);
}
public void publishDirect(List<Object> batch, boolean block) throws InsufficientCapacityException {
int size = 0;
if (batch != null) {
size = batch.size();
if (block) {
int left = size;
int batchIndex = 0;
while (left > 0) {
int availableCapacity = (int) _buffer.remainingCapacity();
if (availableCapacity > 0) {
int n = availableCapacity >= left ? left : availableCapacity;
final long end = _buffer.next(n);
publishBuffer(batch, batchIndex, end - n + 1, end);
left -= n;
batchIndex += n;
} else {
JStormUtils.sleepMs(1);
}
}
} else {
final long end = _buffer.tryNext(size);
publishBuffer(batch, 0, end - size + 1, end);
}
}
}
private void publishBuffer(List<Object> batch, int batchIndex, long bufferBegin, long bufferEnd) {
long at = bufferBegin;
while (at <= bufferEnd) {
final MutableObject m = _buffer.get(at);
m.setObject(batch.get(batchIndex));
at++;
batchIndex++;
}
_buffer.publish(bufferBegin, bufferEnd);
}
public void tryPublish(Object obj) throws InsufficientCapacityException {
publish(obj, false);
}
public void publish(Object obj, boolean block) throws InsufficientCapacityException {
publishDirect(obj, block);
}
protected void publishDirect(Object obj, boolean block) throws InsufficientCapacityException {
final long id;
if (block) {
id = _buffer.next();
} else {
id = _buffer.tryNext(1);
}
final MutableObject m = _buffer.get(id);
m.setObject(obj);
_buffer.publish(id);
}
public void publishCache(Object obj) {
synchronized (_cacheLock) {
_cache.add(obj);
}
}
public int cacheSize() {
synchronized (_cacheLock) {
return _cache.size();
}
}
public void clear() {
while (population() != 0L) {
poll();
}
}
public long population() {
return (writePos() - readPos());
}
public long capacity() {
return _buffer.getBufferSize();
}
public long writePos() {
return _buffer.getCursor();
}
public long readPos() {
return _consumer.get();
}
public float pctFull() {
return (1.0F * population() / capacity());
}
@Override
public Object getState() {
// get readPos then writePos so it's never an under-estimate
long rp = readPos();
long wp = writePos();
state.put("capacity", capacity());
state.put("population", wp - rp);
state.put("write_pos", wp);
state.put("read_pos", rp);
return state;
}
public RingBuffer<MutableObject> get_buffer() {
return _buffer;
}
public SequenceBarrier get_barrier() {
return _barrier;
}
public void publishCallback(Callback cb) {
if (cb != null) {
synchronized (_callbackLock) {
if (_callbacks == null) {
_callbacks = new ArrayList<>();
}
_callbacks.add(cb);
}
}
}
public static class ObjectEventFactory implements EventFactory<MutableObject> {
@Override
public MutableObject newInstance() {
return new MutableObject();
}
}
private class ThreadLocalBatch {
private ArrayList<Object> _batcher;
public ThreadLocalBatch() {
_batcher = new ArrayList<>(_inputBatchSize);
}
public synchronized void addAndFlush(Object obj) {
ArrayList<Object> batchTobeFlushed = add(obj);
if (batchTobeFlushed != null) {
try {
publishDirect(batchTobeFlushed, true);
} catch (InsufficientCapacityException e) {
LOG.warn("Failed to publish batch");
}
}
}
public ArrayList<Object> add(Object obj) {
ArrayList<Object> ret = null;
_batcher.add(obj);
if (_batcher.size() >= _inputBatchSize) {
ret = _batcher;
_batcher = new ArrayList<>(_inputBatchSize);
}
return ret;
}
// May be called by a background thread
public synchronized void flush() {
try {
if (_batcher != null && _batcher.size() > 0) {
publishDirect(_batcher, true);
_batcher = new ArrayList<>(_inputBatchSize);
}
} catch (InsufficientCapacityException e) {
// Ignored we should not block
}
}
}
private class DisruptorFlusher extends Flusher {
private AtomicBoolean _isFlushing = new AtomicBoolean(false);
public DisruptorFlusher(long flushInterval) {
flushIntervalMs = flushInterval;
}
public void run() {
if (_isFlushing.compareAndSet(false, true)) {
_batcher.flush();
_isFlushing.set(false);
}
}
}
}