b1fa6e49af
Add test_module and test_task classes inheriting from respectively task_manager::module and task_manager::task::impl that serve task manager testing.
86 lines
2.1 KiB
C++
86 lines
2.1 KiB
C++
/*
|
|
* Copyright (C) 2022-present ScyllaDB
|
|
*/
|
|
|
|
/*
|
|
* SPDX-License-Identifier: AGPL-3.0-or-later
|
|
*/
|
|
|
|
#ifndef SCYLLA_BUILD_MODE_RELEASE
|
|
|
|
#pragma once
|
|
|
|
#include "task_manager.hh"
|
|
|
|
namespace tasks {
|
|
|
|
class test_module : public task_manager::module {
|
|
private:
|
|
seastar::abort_source _as;
|
|
public:
|
|
test_module(task_manager& tm) noexcept : module(tm, "test") {}
|
|
|
|
seastar::abort_source& abort_source() noexcept override {
|
|
return _as;
|
|
}
|
|
|
|
future<> stop() noexcept override {
|
|
_as.request_abort();
|
|
co_await task_manager::module::stop();
|
|
}
|
|
};
|
|
|
|
class test_task_impl : public task_manager::task::impl {
|
|
private:
|
|
promise<> _finish_run;
|
|
bool _finished = false;
|
|
public:
|
|
test_task_impl(task_manager::module_ptr module, task_id id, uint64_t sequence_number = 0, std::string keyspace = "", std::string table = "", std::string type = "", std::string entity = "", task_id parent_id = task_id::create_null_id())
|
|
: task_manager::task::impl(module, id, sequence_number, std::move(keyspace), std::move(table), std::move(type), std::move(entity), parent_id)
|
|
{}
|
|
|
|
future<> run() override {
|
|
return _finish_run.get_future();
|
|
}
|
|
|
|
friend class test_task;
|
|
};
|
|
|
|
class test_task {
|
|
private:
|
|
task_manager::task_ptr _task;
|
|
public:
|
|
test_task(task_manager::task_ptr task) noexcept : _task(task) {}
|
|
|
|
void finish() noexcept {
|
|
auto& task_impl = dynamic_cast<test_task_impl&>(*_task->_impl);
|
|
if (!task_impl._finished) {
|
|
task_impl._finish_run.set_value();
|
|
task_impl._finished = true;
|
|
}
|
|
}
|
|
|
|
void finish_failed(std::exception_ptr ex) {
|
|
auto& task_impl = dynamic_cast<test_task_impl&>(*_task->_impl);
|
|
if (!task_impl._finished) {
|
|
task_impl._finish_run.set_exception(ex);
|
|
task_impl._finished = true;
|
|
}
|
|
}
|
|
|
|
void register_task() {
|
|
_task->register_task();
|
|
}
|
|
|
|
future<> unregister_task() noexcept {
|
|
auto& task_impl = dynamic_cast<test_task_impl&>(*_task->_impl);
|
|
finish();
|
|
co_await task_impl._done.get_shared_future();
|
|
_task->unregister_task();
|
|
}
|
|
};
|
|
|
|
}
|
|
|
|
#endif
|