iota_surfer/
surf_strategy.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::time::Duration;
6
7use iota_types::{
8    base_types::ObjectRef,
9    transaction::{CallArg, ObjectArg},
10};
11use move_binary_format::normalized;
12use move_core_types::language_storage::StructTag;
13use rand::{Rng, seq::SliceRandom};
14use tokio::time::Instant;
15use tracing::debug;
16
17use crate::surfer_state::{EntryFunction, SurferState};
18
19enum InputObjectPassKind {
20    Value,
21    ByRef,
22    MutRef,
23}
24
25type Type = normalized::Type<normalized::ArcIdentifier>;
26
27#[derive(Clone, Default)]
28pub struct SurfStrategy {
29    min_tx_interval: Duration,
30}
31
32impl SurfStrategy {
33    pub fn new(min_tx_interval: Duration) -> Self {
34        Self { min_tx_interval }
35    }
36
37    /// Given a state and a list of callable Move entry functions,
38    /// explore them for a while, and eventually return. This function may
39    /// not return in some situations, so its important to call it with a
40    /// timeout or select! to ensure the task doesn't block forever.
41    pub async fn surf_for_a_while(
42        &mut self,
43        state: &mut SurferState,
44        mut entry_functions: Vec<EntryFunction>,
45    ) {
46        entry_functions.shuffle(&mut state.rng);
47        for entry in entry_functions {
48            let next_tx_time = Instant::now() + self.min_tx_interval;
49            let Some(args) = Self::choose_function_call_args(state, entry.parameters).await else {
50                debug!(
51                    "Failed to choose arguments for Move function {:?}::{:?}",
52                    entry.module, entry.function
53                );
54                continue;
55            };
56            state
57                .execute_move_transaction(entry.package, entry.module, entry.function, args)
58                .await;
59            tokio::time::sleep_until(next_tx_time).await;
60        }
61    }
62
63    async fn choose_function_call_args(
64        state: &mut SurferState,
65        params: Vec<Type>,
66    ) -> Option<Vec<CallArg>> {
67        let mut args = vec![];
68        let mut chosen_owned_objects = vec![];
69        let mut failed = false;
70        for param in params {
71            let arg = match param {
72                Type::Bool => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<bool>()).unwrap()),
73                Type::U8 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u8>()).unwrap()),
74                Type::U16 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u16>()).unwrap()),
75                Type::U32 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u32>()).unwrap()),
76                Type::U64 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u64>()).unwrap()),
77                Type::U128 => CallArg::Pure(bcs::to_bytes(&state.rng.gen::<u128>()).unwrap()),
78                Type::Address => CallArg::Pure(
79                    bcs::to_bytes(&state.cluster.get_addresses().choose(&mut state.rng)).unwrap(),
80                ),
81                ty @ Type::Datatype(_) => {
82                    match Self::choose_object_call_arg(
83                        state,
84                        InputObjectPassKind::Value,
85                        ty,
86                        &mut chosen_owned_objects,
87                    )
88                    .await
89                    {
90                        Some(arg) => arg,
91                        None => {
92                            failed = true;
93                            break;
94                        }
95                    }
96                }
97                Type::Reference(mut_, ty) => {
98                    let kind = if mut_ {
99                        InputObjectPassKind::MutRef
100                    } else {
101                        InputObjectPassKind::ByRef
102                    };
103                    match Self::choose_object_call_arg(state, kind, *ty, &mut chosen_owned_objects)
104                        .await
105                    {
106                        Some(arg) => arg,
107                        None => {
108                            failed = true;
109                            break;
110                        }
111                    }
112                }
113                Type::U256 | Type::Signer | Type::Vector(_) | Type::TypeParameter(_) => {
114                    failed = true;
115                    break;
116                }
117            };
118            args.push(arg);
119        }
120        if failed {
121            for (struct_tag, obj_ref) in chosen_owned_objects {
122                state
123                    .owned_objects
124                    .get_mut(&struct_tag)
125                    .unwrap()
126                    .insert(obj_ref);
127            }
128            None
129        } else {
130            Some(args)
131        }
132    }
133
134    async fn choose_object_call_arg(
135        state: &mut SurferState,
136        kind: InputObjectPassKind,
137        arg_type: Type,
138        chosen_owned_objects: &mut Vec<(StructTag, ObjectRef)>,
139    ) -> Option<CallArg> {
140        let pool = state.pool.read().await;
141        let type_tag = match arg_type {
142            Type::Datatype(dt) => dt.to_struct_tag(&*pool),
143            _ => {
144                return None;
145            }
146        };
147        drop(pool);
148        let owned = state.matching_owned_objects_count(&type_tag);
149        let shared = state.matching_shared_objects_count(&type_tag).await;
150        let immutable = state.matching_immutable_objects_count(&type_tag).await;
151
152        let total_matching_count = match kind {
153            InputObjectPassKind::Value => owned,
154            InputObjectPassKind::MutRef => owned + shared,
155            InputObjectPassKind::ByRef => owned + shared + immutable,
156        };
157        if total_matching_count == 0 {
158            return None;
159        }
160        let mut n = state.rng.gen_range(0..total_matching_count);
161        if n < owned {
162            let obj_ref = state.choose_nth_owned_object(&type_tag, n);
163            chosen_owned_objects.push((type_tag, obj_ref));
164            return Some(CallArg::Object(ObjectArg::ImmOrOwnedObject(obj_ref)));
165        }
166        n -= owned;
167        if n < shared {
168            let (id, initial_shared_version) = state.choose_nth_shared_object(&type_tag, n).await;
169            return Some(CallArg::Object(ObjectArg::SharedObject {
170                id,
171                initial_shared_version,
172                mutable: matches!(kind, InputObjectPassKind::MutRef),
173            }));
174        }
175        n -= shared;
176        let obj_ref = state.choose_nth_immutable_object(&type_tag, n).await;
177        Some(CallArg::Object(ObjectArg::ImmOrOwnedObject(obj_ref)))
178    }
179}