work_stealing_queue.h 4.97 KB
Newer Older
gejun's avatar
gejun committed
1
// bthread - A M:N threading library to make applications more concurrent.
gejun's avatar
gejun committed
2
// Copyright (c) 2012 Baidu, Inc.
gejun's avatar
gejun committed
3 4 5 6 7 8 9 10 11 12 13 14
// 
// Licensed 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
// 
//     http://www.apache.org/licenses/LICENSE-2.0
// 
// 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.
gejun's avatar
gejun committed
15 16 17 18

// Author: Ge,Jun (gejun@baidu.com)
// Date: Tue Jul 10 17:40:58 CST 2012

gejun's avatar
gejun committed
19 20
#ifndef BTHREAD_WORK_STEALING_QUEUE_H
#define BTHREAD_WORK_STEALING_QUEUE_H
gejun's avatar
gejun committed
21

22 23 24
#include "butil/macros.h"
#include "butil/atomicops.h"
#include "butil/logging.h"
gejun's avatar
gejun committed
25 26 27 28 29 30 31 32 33

namespace bthread {

template <typename T>
class WorkStealingQueue {
public:
    WorkStealingQueue()
        : _bottom(1)
        , _capacity(0)
gejun's avatar
gejun committed
34 35
        , _buffer(NULL)
        , _top(1) {
gejun's avatar
gejun committed
36 37 38 39 40 41 42 43
    }

    ~WorkStealingQueue() {
        delete [] _buffer;
        _buffer = NULL;
    }

    int init(size_t capacity) {
gejun's avatar
gejun committed
44 45
        if (_capacity != 0) {
            LOG(ERROR) << "Already initialized";
gejun's avatar
gejun committed
46 47
            return -1;
        }
gejun's avatar
gejun committed
48 49 50 51 52 53 54
        if (capacity == 0) {
            LOG(ERROR) << "Invalid capacity=" << capacity;
            return -1;
        }
        if (capacity & (capacity - 1)) {
            LOG(ERROR) << "Invalid capacity=" << capacity
                       << " which must be power of 2";
gejun's avatar
gejun committed
55 56 57 58 59 60 61 62 63 64
            return -1;
        }
        _buffer = new(std::nothrow) T[capacity];
        if (NULL == _buffer) {
            return -1;
        }
        _capacity = capacity;
        return 0;
    }

gejun's avatar
gejun committed
65 66 67 68
    // Push an item into the queue.
    // Returns true on pushed.
    // May run in parallel with steal().
    // Never run in parallel with pop() or another push().
gejun's avatar
gejun committed
69
    bool push(const T& x) {
70 71
        const size_t b = _bottom.load(butil::memory_order_relaxed);
        const size_t t = _top.load(butil::memory_order_acquire);
72
        if (b >= t + _capacity) { // Full queue.
gejun's avatar
gejun committed
73 74
            return false;
        }
gejun's avatar
gejun committed
75
        _buffer[b & (_capacity - 1)] = x;
76
        _bottom.store(b + 1, butil::memory_order_release);
gejun's avatar
gejun committed
77 78 79
        return true;
    }

gejun's avatar
gejun committed
80 81 82 83
    // Pop an item from the queue.
    // Returns true on popped and the item is written to `val'.
    // May run in parallel with steal().
    // Never run in parallel with push() or another pop().
gejun's avatar
gejun committed
84
    bool pop(T* val) {
85 86
        const size_t b = _bottom.load(butil::memory_order_relaxed);
        size_t t = _top.load(butil::memory_order_relaxed);
gejun's avatar
gejun committed
87 88 89 90
        if (t >= b) {
            // fast check since we call pop() in each sched.
            // Stale _top which is smaller should not enter this branch.
            return false;
gejun's avatar
gejun committed
91
        }
gejun's avatar
gejun committed
92
        const size_t newb = b - 1;
93 94 95
        _bottom.store(newb, butil::memory_order_relaxed);
        butil::atomic_thread_fence(butil::memory_order_seq_cst);
        t = _top.load(butil::memory_order_relaxed);
gejun's avatar
gejun committed
96
        if (t > newb) {
97
            _bottom.store(b, butil::memory_order_relaxed);
gejun's avatar
gejun committed
98 99 100 101 102 103 104 105
            return false;
        }
        *val = _buffer[newb & (_capacity - 1)];
        if (t != newb) {
            return true;
        }
        // Single last element, compete with steal()
        const bool popped = _top.compare_exchange_strong(
106 107
            t, t + 1, butil::memory_order_seq_cst, butil::memory_order_relaxed);
        _bottom.store(b, butil::memory_order_relaxed);
gejun's avatar
gejun committed
108 109 110
        return popped;
    }

gejun's avatar
gejun committed
111 112 113
    // Steal one item from the queue.
    // Returns true on stolen.
    // May run in parallel with push() pop() or another steal().
gejun's avatar
gejun committed
114
    bool steal(T* val) {
115 116
        size_t t = _top.load(butil::memory_order_acquire);
        size_t b = _bottom.load(butil::memory_order_acquire);
gejun's avatar
gejun committed
117
        if (t >= b) {
gejun's avatar
gejun committed
118
            // Permit false negative for performance considerations.
gejun's avatar
gejun committed
119 120
            return false;
        }
gejun's avatar
gejun committed
121
        do {
122 123
            butil::atomic_thread_fence(butil::memory_order_seq_cst);
            b = _bottom.load(butil::memory_order_acquire);
gejun's avatar
gejun committed
124 125 126 127 128
            if (t >= b) {
                return false;
            }
            *val = _buffer[t & (_capacity - 1)];
        } while (!_top.compare_exchange_strong(t, t + 1,
129 130
                                               butil::memory_order_seq_cst,
                                               butil::memory_order_relaxed));
gejun's avatar
gejun committed
131
        return true;
gejun's avatar
gejun committed
132 133 134
    }

    size_t volatile_size() const {
135 136
        const size_t b = _bottom.load(butil::memory_order_relaxed);
        const size_t t = _top.load(butil::memory_order_relaxed);
gejun's avatar
gejun committed
137
        return (b <= t ? 0 : (b - t));
gejun's avatar
gejun committed
138 139 140 141 142 143 144 145
    }

    size_t capacity() const { return _capacity; }

private:
    // Copying a concurrent structure makes no sense.
    DISALLOW_COPY_AND_ASSIGN(WorkStealingQueue);

146
    butil::atomic<size_t> _bottom;
gejun's avatar
gejun committed
147 148
    size_t _capacity;
    T* _buffer;
149
    butil::atomic<size_t> BAIDU_CACHELINE_ALIGNMENT _top;
gejun's avatar
gejun committed
150 151 152 153
};

}  // namespace bthread

gejun's avatar
gejun committed
154
#endif  // BTHREAD_WORK_STEALING_QUEUE_H