Skip to content

Commit b6537e1

Browse files
snaurykunga
snaury
and
kunga
authored
Add a common TBlockEvents helper class for event interception (#7767)
Co-authored-by: kungurtsev <alexey66rus@gmail.com>
1 parent d356543 commit b6537e1

File tree

5 files changed

+197
-53
lines changed

5 files changed

+197
-53
lines changed

‎ydb/core/testlib/actors/block_events.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
#include "block_events.h"

‎ydb/core/testlib/actors/block_events.h

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
#include "test_runtime.h"
2+
3+
#include <deque>
4+
#include <functional>
5+
6+
namespace NActors {
7+
8+
/**
9+
* Easy blocking for events under the test actor runtime
10+
*
11+
* Matching events are blocked just before they are processed and stashed
12+
* into a deque.
13+
*/
14+
template<class TEvType>
15+
class TBlockEvents : public std::deque<typename TEvType::TPtr> {
16+
public:
17+
TBlockEvents(TTestActorRuntime& runtime, std::function<bool(typename TEvType::TPtr&)> condition = {})
18+
: Runtime(runtime)
19+
, Condition(std::move(condition))
20+
, Holder(Runtime.AddObserver<TEvType>(
21+
[this](typename TEvType::TPtr& ev) {
22+
this->Process(ev);
23+
}))
24+
{}
25+
26+
/**
27+
* Unblocks up to count events at the front of the deque, allowing them
28+
* to be handled by the destination actor.
29+
*/
30+
TBlockEvents& Unblock(size_t count = -1) {
31+
while (!this->empty() && count > 0) {
32+
auto& ev = this->front();
33+
if (!Stopped) {
34+
IEventHandle* ptr = ev.Get();
35+
UnblockedOnce.insert(ptr);
36+
}
37+
ui32 nodeId = ev->GetRecipientRewrite().NodeId();
38+
ui32 nodeIdx = nodeId - Runtime.GetFirstNodeId();
39+
Runtime.Send(ev.Release(), nodeIdx, /* viaActorSystem */ true);
40+
this->pop_front();
41+
--count;
42+
}
43+
return *this;
44+
}
45+
46+
/**
47+
* Stops blocking any new events. Events currently in the deque are
48+
* not unblocked, but may be unblocked at a later time if needed.
49+
*/
50+
TBlockEvents& Stop() {
51+
UnblockedOnce.clear();
52+
Holder.Remove();
53+
Stopped = true;
54+
return *this;
55+
}
56+
57+
private:
58+
void Process(typename TEvType::TPtr& ev) {
59+
IEventHandle* ptr = ev.Get();
60+
auto it = UnblockedOnce.find(ptr);
61+
if (it != UnblockedOnce.end()) {
62+
UnblockedOnce.erase(it);
63+
return;
64+
}
65+
66+
if (Condition && !Condition(ev)) {
67+
return;
68+
}
69+
70+
this->emplace_back(std::move(ev));
71+
}
72+
73+
private:
74+
TTestActorRuntime& Runtime;
75+
std::function<bool(typename TEvType::TPtr&)> Condition;
76+
TTestActorRuntime::TEventObserverHolder Holder;
77+
THashSet<IEventHandle*> UnblockedOnce;
78+
bool Stopped = false;
79+
};
80+
81+
82+
} // namespace NActors

‎ydb/core/testlib/actors/test_runtime_ut.cpp

Lines changed: 110 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include <ydb/core/testlib/actors/test_runtime.h>
2+
#include <ydb/core/testlib/actors/block_events.h>
23
#include <ydb/core/base/appdata.h>
34
#include <ydb/library/actors/core/event_local.h>
45
#include <ydb/library/actors/core/events.h>
@@ -717,6 +718,114 @@ Y_UNIT_TEST_SUITE(TActorTest) {
717718
runtime.WaitFor("value = 42", [&]{ return value == 42; });
718719
UNIT_ASSERT_VALUES_EQUAL(value, 42);
719720
}
720-
};
721+
722+
Y_UNIT_TEST(TestBlockEvents) {
723+
enum EEv {
724+
EvTrigger = EventSpaceBegin(TEvents::ES_PRIVATE)
725+
};
726+
727+
struct TEvTrigger : public TEventLocal<TEvTrigger, EvTrigger> {
728+
int Value;
729+
730+
TEvTrigger(int value)
731+
: Value(value)
732+
{}
733+
};
734+
735+
class TTargetActor : public TActorBootstrapped<TTargetActor> {
736+
public:
737+
TTargetActor(std::vector<int>* ptr)
738+
: Ptr(ptr)
739+
{}
740+
741+
void Bootstrap() {
742+
Become(&TThis::StateWork);
743+
}
744+
745+
private:
746+
STFUNC(StateWork) {
747+
switch (ev->GetTypeRewrite()) {
748+
hFunc(TEvTrigger, Handle);
749+
}
750+
}
751+
752+
void Handle(TEvTrigger::TPtr& ev) {
753+
Ptr->push_back(ev->Get()->Value);
754+
}
755+
756+
private:
757+
std::vector<int>* Ptr;
758+
};
759+
760+
class TSourceActor : public TActorBootstrapped<TSourceActor> {
761+
public:
762+
TSourceActor(const TActorId& target)
763+
: Target(target)
764+
{}
765+
766+
void Bootstrap() {
767+
Become(&TThis::StateWork);
768+
Schedule(TDuration::Seconds(1), new TEvents::TEvWakeup);
769+
}
770+
771+
private:
772+
STFUNC(StateWork) {
773+
switch (ev->GetTypeRewrite()) {
774+
hFunc(TEvents::TEvWakeup, Handle);
775+
}
776+
}
777+
778+
void Handle(TEvents::TEvWakeup::TPtr&) {
779+
Send(Target, new TEvTrigger(++Counter));
780+
Schedule(TDuration::Seconds(1), new TEvents::TEvWakeup);
781+
}
782+
783+
private:
784+
TActorId Target;
785+
int Counter = 0;
786+
};
787+
788+
TTestActorRuntime runtime(2);
789+
runtime.Initialize(MakeEgg());
790+
791+
std::vector<int> values;
792+
auto target = runtime.Register(new TTargetActor(&values), /* nodeIdx */ 1);
793+
auto source = runtime.Register(new TSourceActor(target), /* nodeIdx */ 1);
794+
runtime.EnableScheduleForActor(source);
795+
796+
TBlockEvents<TEvTrigger> block(runtime, [&](TEvTrigger::TPtr& ev){ return ev->GetRecipientRewrite() == target; });
797+
runtime.WaitFor("blocked 3 events", [&]{ return block.size() >= 3; });
798+
UNIT_ASSERT_VALUES_EQUAL(block.size(), 3u);
799+
UNIT_ASSERT_VALUES_EQUAL(values.size(), 0u);
800+
801+
block.Unblock(2);
802+
UNIT_ASSERT_VALUES_EQUAL(block.size(), 1u);
803+
UNIT_ASSERT_VALUES_EQUAL(values.size(), 0u);
804+
805+
runtime.WaitFor("blocked 1 more event", [&]{ return block.size() >= 2; });
806+
UNIT_ASSERT_VALUES_EQUAL(block.size(), 2u);
807+
UNIT_ASSERT_VALUES_EQUAL(values.size(), 2u);
808+
UNIT_ASSERT_VALUES_EQUAL(values.at(0), 1);
809+
UNIT_ASSERT_VALUES_EQUAL(values.at(1), 2);
810+
values.clear();
811+
812+
block.Stop();
813+
runtime.WaitFor("processed 2 more events", [&]{ return values.size() >= 2; });
814+
UNIT_ASSERT_VALUES_EQUAL(block.size(), 2u);
815+
UNIT_ASSERT_VALUES_EQUAL(values.size(), 2u);
816+
UNIT_ASSERT_VALUES_EQUAL(values.at(0), 5);
817+
UNIT_ASSERT_VALUES_EQUAL(values.at(1), 6);
818+
values.clear();
819+
820+
block.Unblock();
821+
UNIT_ASSERT_VALUES_EQUAL(block.size(), 0u);
822+
UNIT_ASSERT_VALUES_EQUAL(values.size(), 0u);
823+
runtime.WaitFor("processed 3 more events", [&]{ return values.size() >= 3; });
824+
UNIT_ASSERT_VALUES_EQUAL(values.size(), 3u);
825+
UNIT_ASSERT_VALUES_EQUAL(values.at(0), 3);
826+
UNIT_ASSERT_VALUES_EQUAL(values.at(1), 4);
827+
UNIT_ASSERT_VALUES_EQUAL(values.at(2), 7);
828+
}
829+
}
721830

722831
}

‎ydb/core/testlib/actors/ya.make

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
LIBRARY()
22

33
SRCS(
4+
block_events.cpp
5+
block_events.h
46
test_runtime.cpp
7+
test_runtime.h
58
)
69

710
PEERDIR(

‎ydb/core/tx/datashard/datashard_ut_read_iterator.cpp

Lines changed: 1 addition & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "read_iterator.h"
55

66
#include <ydb/core/testlib/tablet_helpers.h>
7+
#include <ydb/core/testlib/actors/block_events.h>
78
#include <ydb/core/formats/arrow/arrow_helpers.h>
89
#include <ydb/core/formats/arrow/converter.h>
910
#include <ydb/core/kqp/ut/common/kqp_ut_common.h>
@@ -4627,58 +4628,6 @@ Y_UNIT_TEST_SUITE(DataShardReadIteratorConsistency) {
46274628
"result2: " << result2);
46284629
}
46294630

4630-
template<class TEvType>
4631-
class TBlockEvents : public std::deque<typename TEvType::TPtr> {
4632-
public:
4633-
TBlockEvents(TTestActorRuntime& runtime, std::function<bool(typename TEvType::TPtr&)> condition = {})
4634-
: Runtime(runtime)
4635-
, Condition(std::move(condition))
4636-
, Holder(Runtime.AddObserver<TEvType>(
4637-
[this](typename TEvType::TPtr& ev) {
4638-
this->Process(ev);
4639-
}))
4640-
{}
4641-
4642-
TBlockEvents& Unblock(size_t count = -1) {
4643-
while (!this->empty() && count > 0) {
4644-
auto& ev = this->front();
4645-
IEventHandle* ptr = ev.Get();
4646-
UnblockedOnce.insert(ptr);
4647-
Runtime.Send(ev.Release(), 0, /* viaActorSystem */ true);
4648-
this->pop_front();
4649-
--count;
4650-
}
4651-
return *this;
4652-
}
4653-
4654-
void Stop() {
4655-
UnblockedOnce.clear();
4656-
Holder.Remove();
4657-
}
4658-
4659-
private:
4660-
void Process(typename TEvType::TPtr& ev) {
4661-
IEventHandle* ptr = ev.Get();
4662-
auto it = UnblockedOnce.find(ptr);
4663-
if (it != UnblockedOnce.end()) {
4664-
UnblockedOnce.erase(it);
4665-
return;
4666-
}
4667-
4668-
if (Condition && !Condition(ev)) {
4669-
return;
4670-
}
4671-
4672-
this->emplace_back(std::move(ev));
4673-
}
4674-
4675-
private:
4676-
TTestActorRuntime& Runtime;
4677-
std::function<bool(typename TEvType::TPtr&)> Condition;
4678-
TTestActorRuntime::TEventObserverHolder Holder;
4679-
THashSet<IEventHandle*> UnblockedOnce;
4680-
};
4681-
46824631
Y_UNIT_TEST(Bug_7674_IteratorDuplicateRows) {
46834632
TPortManager pm;
46844633
TServerSettings serverSettings(pm.GetPort(2134));

0 commit comments

Comments
 (0)