work_stealing_queue.h 5.17 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// 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
//
//   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
18 19 20 21 22
// bthread - A M:N threading library to make applications more concurrent.

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

gejun's avatar
gejun committed
23 24
#ifndef BTHREAD_WORK_STEALING_QUEUE_H
#define BTHREAD_WORK_STEALING_QUEUE_H
gejun's avatar
gejun committed
25

26 27 28
#include "butil/macros.h"
#include "butil/atomicops.h"
#include "butil/logging.h"
gejun's avatar
gejun committed
29 30 31 32 33 34 35 36 37

namespace bthread {

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

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

    int init(size_t capacity) {
gejun's avatar
gejun committed
48 49
        if (_capacity != 0) {
            LOG(ERROR) << "Already initialized";
gejun's avatar
gejun committed
50 51
            return -1;
        }
gejun's avatar
gejun committed
52 53 54 55 56 57 58
        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
59 60 61 62 63 64 65 66 67 68
            return -1;
        }
        _buffer = new(std::nothrow) T[capacity];
        if (NULL == _buffer) {
            return -1;
        }
        _capacity = capacity;
        return 0;
    }

gejun's avatar
gejun committed
69 70 71 72
    // 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
73
    bool push(const T& x) {
74 75
        const size_t b = _bottom.load(butil::memory_order_relaxed);
        const size_t t = _top.load(butil::memory_order_acquire);
76
        if (b >= t + _capacity) { // Full queue.
gejun's avatar
gejun committed
77 78
            return false;
        }
gejun's avatar
gejun committed
79
        _buffer[b & (_capacity - 1)] = x;
80
        _bottom.store(b + 1, butil::memory_order_release);
gejun's avatar
gejun committed
81 82 83
        return true;
    }

gejun's avatar
gejun committed
84 85 86 87
    // 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
88
    bool pop(T* val) {
89 90
        const size_t b = _bottom.load(butil::memory_order_relaxed);
        size_t t = _top.load(butil::memory_order_relaxed);
gejun's avatar
gejun committed
91 92 93 94
        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
95
        }
gejun's avatar
gejun committed
96
        const size_t newb = b - 1;
97 98 99
        _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
100
        if (t > newb) {
101
            _bottom.store(b, butil::memory_order_relaxed);
gejun's avatar
gejun committed
102 103 104 105 106 107 108 109
            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(
110 111
            t, t + 1, butil::memory_order_seq_cst, butil::memory_order_relaxed);
        _bottom.store(b, butil::memory_order_relaxed);
gejun's avatar
gejun committed
112 113 114
        return popped;
    }

gejun's avatar
gejun committed
115 116 117
    // 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
118
    bool steal(T* val) {
119 120
        size_t t = _top.load(butil::memory_order_acquire);
        size_t b = _bottom.load(butil::memory_order_acquire);
gejun's avatar
gejun committed
121
        if (t >= b) {
gejun's avatar
gejun committed
122
            // Permit false negative for performance considerations.
gejun's avatar
gejun committed
123 124
            return false;
        }
gejun's avatar
gejun committed
125
        do {
126 127
            butil::atomic_thread_fence(butil::memory_order_seq_cst);
            b = _bottom.load(butil::memory_order_acquire);
gejun's avatar
gejun committed
128 129 130 131 132
            if (t >= b) {
                return false;
            }
            *val = _buffer[t & (_capacity - 1)];
        } while (!_top.compare_exchange_strong(t, t + 1,
133 134
                                               butil::memory_order_seq_cst,
                                               butil::memory_order_relaxed));
gejun's avatar
gejun committed
135
        return true;
gejun's avatar
gejun committed
136 137 138
    }

    size_t volatile_size() const {
139 140
        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
141
        return (b <= t ? 0 : (b - t));
gejun's avatar
gejun committed
142 143 144 145 146 147 148 149
    }

    size_t capacity() const { return _capacity; }

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

150
    butil::atomic<size_t> _bottom;
gejun's avatar
gejun committed
151 152
    size_t _capacity;
    T* _buffer;
153
    butil::atomic<size_t> BAIDU_CACHELINE_ALIGNMENT _top;
gejun's avatar
gejun committed
154 155 156 157
};

}  // namespace bthread

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