feat(resharding): introduce resharding actor #12217

Trisfald · 2024-10-14T13:23:24Z

Main changes introduced by this PR:

Calling FlatStorageResharder resume and start in all code paths.
Added a very thin ReshardingActor and related Sender/Request items. Unfortunately, this meant propagating an extra argument through Client and Chain.
Moved FlatStorageResharder in ReshardingManager. Wrapped inside Option because it makes initialization of Chain easier, and it's not needed in many tests anyway.
Refactored FlatStorageResharder to work with a Sender. Functionality is the same as before.
A lot of uninteresting changes just to make old tests compile.

…ding

…ctor

…ding-actor

wacban

LGTM

wacban · 2024-10-15T10:37:51Z

chain/chain/src/chain.rs

@@ -361,7 +362,9 @@ impl Chain {
        let resharding_manager = ReshardingManager::new(


Do we even need resharding manager in the view client?

I don't see a reason. I'll make a note to ask @shreyan-gupta and then remove the manager from there

Shouldn't be needed, but we would have to make it Option then?

wacban · 2024-10-15T10:40:12Z

chain/chain/src/flat_storage_resharder.rs

@@ -52,13 +54,24 @@ use near_store::{ShardUId, StorageError};
 pub struct FlatStorageResharder {
    runtime: Arc<dyn RuntimeAdapter>,
    resharding_event: Arc<Mutex<Option<FlatStorageReshardingEventStatus>>>,
+    scheduler: messaging::Sender<FlatStorageSplitShardRequest>,


nit: Maybe import messaging? I don't think there is any ambiguity about it.

wacban · 2024-10-15T10:40:42Z

chain/chain/src/flat_storage_resharder.rs

+    /// * `runtime`: runtime adapter
+    /// * `scheduler`: component used to schedule the background tasks
+    /// * `controller`: manages the execution of the background tasks


That's pretty!

wacban · 2024-10-15T10:43:21Z

chain/chain/src/flat_storage_resharder.rs

-fn split_shard_task(resharder: FlatStorageResharder, controller: FlatStorageResharderController) {
-    let task_status = split_shard_task_impl(resharder.clone(), controller.clone());
+pub fn split_shard_task(resharder: FlatStorageResharder) {
+    let task_status = split_shard_task_impl(resharder.clone());


Sanity check: What's up with the clones?

Good catch, it's not really needed anymore. Reference is enough.

wacban · 2024-10-15T10:44:57Z

chain/chain/src/flat_storage_resharder.rs

@@ -501,7 +501,7 @@ pub enum FlatStorageReshardingTaskStatus {

 /// Helps control the flat storage resharder operation. More specifically,
 /// it has a way to know when the background task is done or to interrupt it.
-#[derive(Clone)]
+#[derive(Clone, Debug)]


I'm surprised that sender and receiver implement Debug but alright.

Caught me by surprise as well

wacban · 2024-10-15T10:52:16Z

chain/chain/src/resharding/manager.rs

 }

 impl ReshardingManager {
    pub fn new(
        store: Store,
        epoch_manager: Arc<dyn EpochManagerAdapter>,
+        runtime_adapter: Arc<dyn RuntimeAdapter>,


Just for my education, why do you need the runtime adapter in this PR?

The reason is that now FlatStorageResharder is owned and created by ReshardingManager. runtime is needed by FlatStorageResharder

wacban · 2024-10-15T10:53:00Z

chain/chain/src/resharding/manager.rs

+                ReshardingEventType::SplitShard(split_shard_event.clone()),
+                &next_shard_layout,
+            )?,
+            None => tracing::info!(target: "resharding", "flat storage resharder not initialized"),


should that be tracing::error?

wacban · 2024-10-15T10:53:46Z

chain/chain/src/resharding/resharding_actor.rs

+    }
+
+    pub fn handle_flat_storage_split_shard_request(&mut self, msg: FlatStorageSplitShardRequest) {
+        split_shard_task(msg.resharder);


nit: Would it be any nicer as a method? msg.resharder.split_shard_task() ?

Yes, I did the change and it looks nicer to me

wacban · 2024-10-15T10:54:16Z

chain/chain/src/resharding/types.rs

+
+#[derive(actix::Message, Clone, Debug)]
+#[rtype(result = "()")]
+pub struct FlatStorageSplitShardRequest {


Can you add comments to the pub structs?

wacban · 2024-10-15T10:55:06Z

chain/chain/src/resharding/types.rs

+
+#[derive(Clone, near_async::MultiSend, near_async::MultiSenderFrom)]
+pub struct ReshardingSender {
+    pub flat_storage_split_shard_send: Sender<FlatStorageSplitShardRequest>,


In the future we'd want a sender for any type of resharding request and SplitShard would be one of the variants in the Request enum. No need to do it now though.

…ding-actor

tayfunelmas · 2024-10-15T14:39:41Z

chain/chain/src/resharding/manager.rs

+                ReshardingEventType::SplitShard(split_shard_event.clone()),
+                &next_shard_layout,
+            )?,
+            None => tracing::error!(target: "resharding", "flat storage resharder not initialized"),


should not we panic if the resharder is not initialized? or it is based on some protocol versioning? if so, it should not error message right?

Converting resharding sender to non-option should remove this case.

tayfunelmas · 2024-10-15T14:43:08Z

chain/chain/src/flat_storage_resharder.rs

@@ -609,7 +605,10 @@ mod tests {
    }

    /// Generic test setup.
-    fn create_fs_resharder(shard_layout: ShardLayout) -> (Chain, FlatStorageResharder) {
+    fn create_fs_resharder(


nit. what does fs stand for here? if flat_storage, can we expand the name?

it stands for flat storage yes

function name aged quite badly, I'll change it!

tayfunelmas · 2024-10-15T14:45:49Z

chain/chain/src/resharding/manager.rs

@@ -23,15 +27,27 @@ pub struct ReshardingManager {
    /// A handle that allows the main process to interrupt resharding if needed.
    /// This typically happens when the main process is interrupted.
    pub resharding_handle: ReshardingHandle,
+    /// Takes care of performing resharding on the flat storage.
+    pub flat_storage_resharder: Option<FlatStorageResharder>,


I would prefer having this non-optional (maybe add a TODO as it is now convenient for this PR) in the future since as far as I know, the resharding will fail if this is not initialized yet? or can it still succeed?

Would too prefer to have this non-optional. I think we should address this as part of this PR itself. For the places where we don't have a sender or don't need to pass a sender, we can just pass in noop() sender.

I'll try to remove the Option and pass a noop()!

codecov · 2024-10-15T15:01:14Z

Codecov Report

Attention: Patch coverage is 89.74359% with 24 lines in your changes missing coverage. Please review.

Project coverage is 71.72%. Comparing base (3cb74c2) to head (93b7556).
Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
chain/chain/src/flat_storage_resharder.rs	89.50%	12 Missing and 5 partials ⚠️
chain/chain/src/flat_storage_creator.rs	33.33%	2 Missing ⚠️
tools/database/src/resharding_v2.rs	0.00%	2 Missing ⚠️
chain/chain/src/resharding/manager.rs	92.30%	0 Missing and 1 partial ⚠️
chain/client/src/test_utils/test_env.rs	0.00%	1 Missing ⚠️
tools/speedy_sync/src/main.rs	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #12217   +/-   ##
=======================================
  Coverage   71.72%   71.72%           
=======================================
  Files         833      835    +2     
  Lines      166700   166684   -16     
  Branches   166700   166684   -16     
=======================================
- Hits       119573   119562   -11     
+ Misses      41902    41901    -1     
+ Partials     5225     5221    -4

Flag	Coverage Δ
backward-compatibility	`0.17% <0.00%> (-0.01%)`	⬇️
db-migration	`0.17% <0.00%> (-0.01%)`	⬇️
genesis-check	`1.25% <0.00%> (-0.01%)`	⬇️
integration-tests	`38.88% <58.54%> (+0.15%)`	⬆️
linux	`71.39% <83.33%> (-0.02%)`	⬇️
linux-nightly	`71.32% <89.74%> (+<0.01%)`	⬆️
macos	`53.82% <82.96%> (-0.50%)`	⬇️
pytests	`1.57% <0.00%> (-0.01%)`	⬇️
sanity-checks	`1.37% <0.00%> (-0.01%)`	⬇️
unittests	`65.52% <82.96%> (-0.03%)`	⬇️
upgradability	`0.21% <0.00%> (-0.01%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

shreyan-gupta

Looks good overall but have a couple of comments

shreyan-gupta · 2024-10-15T15:58:10Z

chain/chain/src/flat_storage_resharder.rs

        let resharder = self.clone();
-        let task = Box::new(move || split_shard_task(resharder, controller));
-        scheduler.schedule(task);
+        self.scheduler.send(FlatStorageSplitShardRequest { resharder });


This does not look right, why are we sending an instance of resharder along with request? Can the actor hold an instance of resharder as a member instead?

Yes, it should be easy to make the actor hold copy of resharder

shreyan-gupta · 2024-10-15T16:12:37Z

chain/client/src/test_utils/mock_resharding_sender.rs

+    split_shard_request: Arc<RwLock<VecDeque<FlatStorageSplitShardRequest>>>,
+}
+
+impl CanSend<FlatStorageSplitShardRequest> for MockReshardingSender {


This is the old framework for tests, before test loop. Let's try to keep it as lean as possible and not introduce any new functionality here? It would be really really really really really really great it we can just have a noop sender and kill all this code and just say we don't support resharding for old framework.

I'll try to kill this code and replace with the least amount of code just to make stuff compile

shreyan-gupta · 2024-10-15T16:15:31Z

chain/chain/src/resharding/manager.rs

@@ -23,15 +27,27 @@ pub struct ReshardingManager {
    /// A handle that allows the main process to interrupt resharding if needed.
    /// This typically happens when the main process is interrupted.
    pub resharding_handle: ReshardingHandle,
+    /// Takes care of performing resharding on the flat storage.
+    pub flat_storage_resharder: Option<FlatStorageResharder>,


Would too prefer to have this non-optional. I think we should address this as part of this PR itself. For the places where we don't have a sender or don't need to pass a sender, we can just pass in noop() sender.

shreyan-gupta · 2024-10-15T16:16:48Z

chain/chain/src/resharding/manager.rs

+                ReshardingEventType::SplitShard(split_shard_event.clone()),
+                &next_shard_layout,
+            )?,
+            None => tracing::error!(target: "resharding", "flat storage resharder not initialized"),


Converting resharding sender to non-option should remove this case.

shreyan-gupta · 2024-10-15T16:29:06Z

chain/chain/src/resharding/resharding_actor.rs

+    }
+
+    pub fn handle_flat_storage_split_shard_request(&mut self, msg: FlatStorageSplitShardRequest) {
+        msg.resharder.split_shard_task();


I'm honestly not too happy that we are calling a function from resharder here.

For the purposes of reducing mental burden of identifying what happens where, ideally all the logic to deal with the heavy task of splitting should have been a part of the ReshardingActor class. This is too much of indirection to deal with.

Is it possible for us to push the logic of splitting the flat storage here?

I thought the general sentiment was to keep the actor as simple as possible, thus all the work and unit testing is done in the resharder. Any other thoughts?

Would love input from other folks, @wacban @Longarithm!

I don't have strong feelings either way. My only recommendation would be to keep the actor api part (typically the handle method) as slim as possible. This way we can properly unit test the actual implementation without needed to spin up the actor or worry about async. So something like this:

fn handle(request) { the_actual_implementation(request); } mod tests { // here the the_actual_implementation can be unit tested without needing to bootstrap an entire actor

shreyan-gupta · 2024-10-15T16:33:17Z

nearcore/src/lib.rs

@@ -401,6 +405,7 @@ pub fn start_with_config_and_synchronization(
        partial_witness_actor.clone().with_auto_span_context().into_multi_sender(),
        true,
        None,
+        resharding_sender.into_multi_sender(),


Ugh, we need a better way to handle adding new actors and new actix requests. Note to self; I'll see if I can come up with better ways of handling all actor senders in a single unified struct.

shreyan-gupta · 2024-10-15T16:48:35Z

chain/chain/src/resharding/manager.rs

+            FlatStorageResharder::new(
+                runtime_adapter,
+                sender.into_sender(),
+                FlatStorageResharderController::from_resharding_handle(resharding_handle.clone()),


I'm still unsure how does the controller work and what do we do with the sender and receiver here? I'll have to go back and do a quick read up on how crossbeam_channels work but we should ideally not be introducing a new framework here given we have actors and message senders? Thoughts?

I planned to remove the channels in another PR because I thought they could come in handy but they did not. I can push the change here

…ding-actor

Trisfald · 2024-10-16T10:57:11Z

I've addressed most of the comments!

I left out two things:

Making the ReshardingActor hold a resharder instance: it's a bigger change than expected because it introduces a circular dependency between the two
Merging ReshardingActor with FlatStorageResharder: resharder is used both by the actor and by the logic which handles flat state status at startup, so it is convenient to have this separation at the moment

I plan to do further refactoring and improvements in separate PRs

…ding-actor

Trisfald added 30 commits September 19, 2024 10:59

fix typo

51f9f44

skeleton for flat storage resharding V3 update

eb421de

skeleton of flat storage resharder

c1a71a6

add TODO details

9234160

add API stub to resume or start resharding of flat storage

c39b476

flat storage resharding start fn

f977326

many improvements

d6823d3

add test

0195907

improve resharder api

c7a6875

Merge remote-tracking branch 'origin/master' into flat-storage-reshar…

4c7e9b2

…ding

use the new flat store adapter

79730a4

implement simple key value copy

ad83e0b

add implemnetation of flat storage resharding interrupt

ae82453

change tests setup to use a test chain

e735755

handle more DB columns

fc38090

improve comments

75c42b5

Merge remote-tracking branch 'origin/master' into flat-storage-reshar…

3848c40

…ding

fix unit tests

306c502

add parent shard deletion in task

b9f11d4

update protocol schema

7fe014e

remove todo

d02fc52

remove parent flat storage through the manager

cf8ec95

multi line imports

f45cc9e

apply code review suggestions

3c188c5

add note about parent shard flat storage invariant

ac6d642

add test reject_split_shard_if_parent_is_not_ready

9269c37

use a dedicated type to represent task statuses

7ad602e

refactor copy_kv_to_child into a free function

d0b3125

address code review comments

df869fc

refactor ReshardingEventType

e1f3b5d

Trisfald added 6 commits October 14, 2024 11:29

propagate error

17538f9

make get_parent_shard_and_status a resharder function

ebd144a

Merge branch 'flat-storage-resharding' into flat-storage-resharding-a…

534d9bc

…ctor

Merge remote-tracking branch 'origin/master' into flat-storage-reshar…

879b1b5

…ding-actor

fix merge

794669f

fix references to last block of old shard layout

f2ee9d3

Trisfald requested a review from a team as a code owner October 14, 2024 13:23

Trisfald requested review from tayfunelmas, wacban and shreyan-gupta October 14, 2024 13:23

wacban approved these changes Oct 15, 2024

View reviewed changes

Trisfald added 4 commits October 15, 2024 14:25

feedback

318957e

Merge remote-tracking branch 'origin/master' into flat-storage-reshar…

b8a5a2d

…ding-actor

make split_shard_task a member of FlatStorageResharder

c5e2886

prevent double deletion of parent flat storage

465706a

tayfunelmas approved these changes Oct 15, 2024

View reviewed changes

tayfunelmas reviewed Oct 15, 2024

View reviewed changes

shreyan-gupta reviewed Oct 15, 2024

View reviewed changes

Trisfald added 5 commits October 16, 2024 10:45

rename create_fs_resharder()

e51a012

remove channels from flat storage resharder

232b89b

remove mock_resharding_sender

5d042fc

make FlatStorageResharder not an Option in manager

d817ebb

Merge remote-tracking branch 'origin/master' into flat-storage-reshar…

3464767

…ding-actor

Merge remote-tracking branch 'origin/master' into flat-storage-reshar…

93b7556

…ding-actor

Trisfald added this pull request to the merge queue Oct 16, 2024

Merged via the queue into near:master with commit aa48e52 Oct 16, 2024
24 of 25 checks passed

Trisfald deleted the flat-storage-resharding-actor branch October 16, 2024 20:22

		@@ -361,7 +362,9 @@ impl Chain {
		let resharding_manager = ReshardingManager::new(

feat(resharding): introduce resharding actor #12217

feat(resharding): introduce resharding actor #12217

Conversation

Trisfald commented Oct 14, 2024

wacban left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shreyan-gupta Oct 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Oct 15, 2024 • edited Loading

Codecov Report

shreyan-gupta left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Trisfald commented Oct 16, 2024

shreyan-gupta Oct 15, 2024 •

edited

Loading

codecov bot commented Oct 15, 2024 •

edited

Loading