gate.hh

/*
 * This file is open source software, licensed to you under the terms
 * of the Apache License, Version 2.0 (the "License").  See the NOTICE file
 * distributed with this work for additional information regarding copyright
 * ownership.  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.
 */
/*
 * Copyright 2014 Cloudius Systems
 */
 
#pragma once
 
#include <seastar/core/future.hh>
#include <seastar/util/std-compat.hh>
#include <exception>
 
namespace seastar {
 
 
class gate_closed_exception : public std::exception {
public:
    virtual const char* what() const noexcept override {
        return "gate closed";
    }
};
 
class gate {
    size_t _count = 0;
    std::optional<promise<>> _stopped;
public:
    gate() = default;
    gate(const gate&) = delete;
    gate(gate&&) = default;
    gate& operator=(gate&&) = default;
    ~gate() {
        assert(!_count && "gate destroyed with outstanding requests");
    }
    bool try_enter() noexcept {
        bool opened = !_stopped;
        if (opened) {
            ++_count;
        }
        return opened;
    }
    void enter() {
        if (!try_enter()) {
            throw gate_closed_exception();
        }
    }
    void leave() noexcept {
        --_count;
        if (!_count && _stopped) {
            _stopped->set_value();
        }
    }
    void check() {
        if (_stopped) {
            throw gate_closed_exception();
        }
    }
    future<> close() noexcept {
        assert(!_stopped && "seastar::gate::close() cannot be called more than once");
        _stopped = std::make_optional(promise<>());
        if (!_count) {
            _stopped->set_value();
        }
        return _stopped->get_future();
    }
 
    size_t get_count() const noexcept {
        return _count;
    }
 
    bool is_closed() const noexcept {
        return bool(_stopped);
    }
 
    class holder {
        gate* _g;
 
    public:
        holder() noexcept : _g(nullptr) { }
 
        explicit holder(gate& g) : _g(&g) {
            _g->enter();
        }
 
        holder(const holder& x) noexcept : _g(x._g) {
            if (_g) {
                _g->_count++;
            }
        }
 
        holder(holder&& x) noexcept : _g(std::exchange(x._g, nullptr)) { }
 
        ~holder() {
            release();
        }
 
        holder& operator=(const holder& x) noexcept {
            if (x._g != _g) {
                release();
                _g = x._g;
                if (_g) {
                    _g->_count++;
                }
            }
            return *this;
        }
 
        holder& operator=(holder&& x) noexcept {
            if (&x != this) {
                release();
                _g = std::exchange(x._g, nullptr);
            }
            return *this;
        }
 
        void release() noexcept {
            if (_g) {
                _g->leave();
                _g = nullptr;
            }
        }
    };
 
    holder hold() {
        return holder(*this);
    }
};
 
namespace internal {
 
template <typename Func>
inline
auto
invoke_func_with_gate(gate& g, Func&& func) noexcept {
    return futurize_invoke(std::forward<Func>(func)).finally([&g] { g.leave(); });
}
 
} // namespace intgernal
 
template <typename Func>
inline
auto
with_gate(gate& g, Func&& func) {
    g.enter();
    return internal::invoke_func_with_gate(g, std::forward<Func>(func));
}
 
template <typename Func>
inline
auto
try_with_gate(gate& g, Func&& func) noexcept {
    if (!g.try_enter()) {
        using futurator = futurize<std::result_of_t<Func()>>;
        return futurator::make_exception_future(gate_closed_exception());
    }
    return internal::invoke_func_with_gate(g, std::forward<Func>(func));
}
 
}