001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.camel.util.concurrent;
018
019import java.util.PriorityQueue;
020import java.util.concurrent.Executor;
021import java.util.concurrent.TimeUnit;
022import java.util.concurrent.atomic.AtomicInteger;
023import java.util.concurrent.locks.Condition;
024import java.util.concurrent.locks.ReentrantLock;
025import java.util.function.Consumer;
026
027/**
028 * A completion service that orders the completed tasks in the same order as they where submitted.
029 */
030public class AsyncCompletionService<V> {
031
032    private final Executor executor;
033    private final boolean ordered;
034    private final PriorityQueue<Task> queue;
035    private final AtomicInteger nextId = new AtomicInteger();
036    private final AtomicInteger index = new AtomicInteger();
037    private final ReentrantLock lock;
038    private final Condition available;
039
040    public AsyncCompletionService(Executor executor, boolean ordered) {
041        this(executor, ordered, null, 0);
042    }
043
044    public AsyncCompletionService(Executor executor, boolean ordered, ReentrantLock lock) {
045        this(executor, ordered, lock, 0);
046    }
047
048    public AsyncCompletionService(Executor executor, boolean ordered, ReentrantLock lock, int capacity) {
049        this.executor = executor;
050        this.ordered = ordered;
051        this.lock = lock != null ? lock : new ReentrantLock();
052        this.available = this.lock.newCondition();
053        if (capacity > 0) {
054            queue = new PriorityQueue<>(capacity);
055        } else {
056            queue = new PriorityQueue<>();
057        }
058    }
059
060    public ReentrantLock getLock() {
061        return lock;
062    }
063
064    public void submit(Consumer<Consumer<V>> runner) {
065        Task f = new Task(nextId.getAndIncrement(), runner);
066        this.executor.execute(f);
067    }
068
069    public void skip() {
070        index.incrementAndGet();
071    }
072
073    public V pollUnordered() {
074        final ReentrantLock lock = this.lock;
075        lock.lock();
076        try {
077            Task t = queue.poll();
078            return t != null ? t.result : null;
079        } finally {
080            lock.unlock();
081        }
082    }
083
084    public V poll() {
085        final ReentrantLock lock = this.lock;
086        lock.lock();
087        try {
088            Task t = queue.peek();
089            if (t != null && (!ordered || index.compareAndSet(t.id, t.id + 1))) {
090                queue.poll();
091                return t.result;
092            } else {
093                return null;
094            }
095        } finally {
096            lock.unlock();
097        }
098    }
099
100    public V poll(long timeout, TimeUnit unit) throws InterruptedException {
101        long nanos = unit.toNanos(timeout);
102        final ReentrantLock lock = this.lock;
103        lock.lockInterruptibly();
104        try {
105            for (;;) {
106                Task t = queue.peek();
107                if (t != null && (!ordered || index.compareAndSet(t.id, t.id + 1))) {
108                    queue.poll();
109                    return t.result;
110                }
111                if (nanos <= 0) {
112                    return null;
113                } else {
114                    nanos = available.awaitNanos(nanos);
115                }
116            }
117        } finally {
118            lock.unlock();
119        }
120    }
121
122    public V take() throws InterruptedException {
123        final ReentrantLock lock = this.lock;
124        lock.lockInterruptibly();
125        try {
126            for (;;) {
127                Task t = queue.peek();
128                if (t != null && (!ordered || index.compareAndSet(t.id, t.id + 1))) {
129                    queue.poll();
130                    return t.result;
131                }
132                available.await();
133            }
134        } finally {
135            lock.unlock();
136        }
137    }
138
139    private void complete(Task task) {
140        final ReentrantLock lock = this.lock;
141        lock.lock();
142        try {
143            queue.add(task);
144            available.signalAll();
145        } finally {
146            lock.unlock();
147        }
148    }
149
150    private class Task implements Runnable, Comparable<Task>, Consumer<V> {
151        private final int id;
152        private final Consumer<Consumer<V>> runner;
153        private V result;
154
155        Task(int id, Consumer<Consumer<V>> runner) {
156            this.id = id;
157            this.runner = runner;
158        }
159
160        @Override
161        public void run() {
162            runner.accept(this);
163        }
164
165        @Override
166        public void accept(V result) {
167            this.result = result;
168            complete(this);
169        }
170
171        @Override
172        public int compareTo(Task other) {
173            return Integer.compare(this.id, other.id);
174        }
175
176        @Override
177        public String toString() {
178            return "SubmitOrderedTask[" + this.id + "]";
179        }
180    }
181}