@@ -32,9 +32,9 @@ use crate::store::repositories::UndecidedBlocksRepository;
3232/// Handles the `RestreamProposal` message from the consensus engine.
3333///
3434/// This is called when the consensus engine requests to restream a proposal for a specific height and round.
35- /// If a valid round is provided, the application first fetches the block from the valid round,
36- /// updates its round to the new round, and stores it. Then, it fetches the
37- /// block for the specified height and round, streams the proposal parts, and sends them over the network .
35+ /// The block is looked up by height and block hash (ignoring round), so it will be found
36+ /// regardless of which round it was originally stored under. The block's round and valid_round
37+ /// are updated to match the current proposal context before restreaming .
3838///
3939/// ## Errors
4040/// - If no block is found for the specified height and round
@@ -107,18 +107,17 @@ async fn get_block_to_restream(
107107 valid_round : Round ,
108108 block_hash : BlockHash ,
109109) -> eyre:: Result < Option < ConsensusBlock > > {
110- if valid_round. is_defined ( )
111- && let Some ( mut block) = undecided_blocks
112- . get ( height, valid_round, block_hash)
113- . await
114- . wrap_err_with ( || {
115- format ! (
116- "Failed to fetch block from valid round for restreaming it \
117- (height={height}, valid_round={valid_round}, block_hash={block_hash})"
118- )
119- } ) ?
120- {
121- // Update the block for the new round and store it
110+ let block = undecided_blocks
111+ . get_first ( height, block_hash)
112+ . await
113+ . wrap_err_with ( || {
114+ format ! (
115+ "Failed to fetch block for restreaming \
116+ (height={height}, round={round}, block_hash={block_hash})"
117+ )
118+ } ) ?;
119+
120+ if let Some ( mut block) = block {
122121 block. round = round;
123122 block. valid_round = valid_round;
124123
@@ -127,24 +126,14 @@ async fn get_block_to_restream(
127126 . await
128127 . wrap_err_with ( || {
129128 format ! (
130- "Failed to store updated undecided block from valid round before restreaming it \
131- (height={height}, valid_round={valid_round}, block_hash={block_hash})"
132- )
133- } ) ?;
134-
135- Ok ( Some ( block) )
136- } else {
137- let block_to_restream = undecided_blocks
138- . get ( height, round, block_hash)
139- . await
140- . wrap_err_with ( || {
141- format ! (
142- "Failed to fetch block from round for restreaming it \
129+ "Failed to store updated block before restreaming \
143130 (height={height}, round={round}, block_hash={block_hash})"
144131 )
145132 } ) ?;
146133
147- Ok ( block_to_restream)
134+ Ok ( Some ( block) )
135+ } else {
136+ Ok ( None )
148137 }
149138}
150139
@@ -179,25 +168,22 @@ mod tests {
179168 }
180169
181170 #[ tokio:: test]
182- async fn get_block_from_valid_round_and_store_update ( ) {
171+ async fn get_block_found_and_updated ( ) {
183172 let mut mock_repo = MockUndecidedBlocksRepository :: new ( ) ;
184173
185174 let height = Height :: new ( 10 ) ;
186- let round = Round :: new ( 5 ) ; // The new round we are proposing in
187- let valid_round = Round :: new ( 3 ) ; // The previous round where we saw the block
175+ let round = Round :: new ( 5 ) ;
176+ let valid_round = Round :: new ( 3 ) ;
188177 let block_hash = BlockHash :: default ( ) ;
189178
190- // Original block proposed in round 3
191- let original_block = create_dummy_block ( height, valid_round, Round :: Nil ) ;
179+ let original_block = create_dummy_block ( height, Round :: new ( 0 ) , Round :: Nil ) ;
192180
193- // Expectation: Fetch from valid_round (3)
194181 mock_repo
195- . expect_get ( )
196- . with ( eq ( height) , eq ( valid_round ) , eq ( block_hash) )
182+ . expect_get_first ( )
183+ . with ( eq ( height) , eq ( block_hash) )
197184 . times ( 1 )
198- . returning ( move |_, _, _ | Ok ( Some ( original_block. clone ( ) ) ) ) ;
185+ . returning ( move |_, _| Ok ( Some ( original_block. clone ( ) ) ) ) ;
199186
200- // Expectation: Store the block updated with the NEW round (5) and valid_round (3)
201187 mock_repo
202188 . expect_store ( )
203189 . withf ( move |b| b. round == round && b. valid_round == valid_round)
@@ -207,79 +193,30 @@ mod tests {
207193 let result =
208194 get_block_to_restream ( & mock_repo, height, round, valid_round, block_hash) . await ;
209195
210- assert ! ( result. is_ok( ) ) ;
211- let block = result. unwrap ( ) ;
212- assert ! ( block. is_some( ) ) ;
213-
214- let b = block. unwrap ( ) ;
215- assert_eq ! ( b. round, round) ; // Ensure returned block has updated round
216- assert_eq ! ( b. valid_round, valid_round) ;
217- }
218-
219- #[ tokio:: test]
220- async fn get_block_no_valid_round_fetch_current ( ) {
221- let mut mock_repo = MockUndecidedBlocksRepository :: new ( ) ;
222-
223- let height = Height :: new ( 10 ) ;
224- let round = Round :: new ( 5 ) ;
225- let valid_round = Round :: Nil ; // No valid round
226- let block_hash = BlockHash :: default ( ) ;
227-
228- // A block proposed for the first time at round 5, no valid round
229- let current_block = create_dummy_block ( height, round, valid_round) ;
230-
231- // Expectation: we are restreaming this block because we received it at round 5.
232- mock_repo
233- . expect_get ( )
234- . with ( eq ( height) , eq ( round) , eq ( block_hash) )
235- . times ( 1 )
236- . returning ( move |_, _, _| Ok ( Some ( current_block. clone ( ) ) ) ) ;
237-
238- // Expectation: Store should NOT be called
239- mock_repo. expect_store ( ) . times ( 0 ) ;
240-
241- let result =
242- get_block_to_restream ( & mock_repo, height, round, valid_round, block_hash) . await ;
243-
244- assert ! ( result. is_ok( ) ) ;
245- assert_eq ! ( result. unwrap( ) . unwrap( ) . round, round) ;
196+ let block = result. unwrap ( ) . expect ( "block should be found" ) ;
197+ assert_eq ! ( block. round, round) ;
198+ assert_eq ! ( block. valid_round, valid_round) ;
246199 }
247200
248201 #[ tokio:: test]
249- async fn fallback_when_valid_round_block_missing ( ) {
202+ async fn get_block_not_found ( ) {
250203 let mut mock_repo = MockUndecidedBlocksRepository :: new ( ) ;
251204
252205 let height = Height :: new ( 10 ) ;
253206 let round = Round :: new ( 5 ) ;
254207 let valid_round = Round :: new ( 3 ) ;
255208 let block_hash = BlockHash :: default ( ) ;
256209
257- let block_at_current = create_dummy_block ( height, round, valid_round) ;
258-
259- // 1. First fetch at valid_round returns None
260- // The proposed value can be valid for the proposer but not for us.
261- mock_repo
262- . expect_get ( )
263- . with ( eq ( height) , eq ( valid_round) , eq ( block_hash) )
264- . times ( 1 )
265- . returning ( |_, _, _| Ok ( None ) ) ;
266-
267- // 2. Fallback: fetch at current round
268- // Since it was restreamed by the proposer, we have it as current's round value.
269210 mock_repo
270- . expect_get ( )
271- . with ( eq ( height) , eq ( round ) , eq ( block_hash) )
211+ . expect_get_first ( )
212+ . with ( eq ( height) , eq ( block_hash) )
272213 . times ( 1 )
273- . returning ( move |_, _, _| Ok ( Some ( block_at_current. clone ( ) ) ) ) ;
274-
275- // Store should NOT be called in the fallback path
276- mock_repo. expect_store ( ) . times ( 0 ) ;
214+ . returning ( |_, _| Ok ( None ) ) ;
277215
278216 let result =
279217 get_block_to_restream ( & mock_repo, height, round, valid_round, block_hash) . await ;
280218
281- assert ! ( result. is_ok( ) ) ;
282- assert ! ( result. unwrap( ) . is_some( ) ) ;
219+ assert ! ( result. unwrap( ) . is_none( ) ) ;
283220 }
284221
285222 #[ tokio:: test]
@@ -290,8 +227,8 @@ mod tests {
290227 let valid_round = Round :: Nil ;
291228
292229 mock_repo
293- . expect_get ( )
294- . returning ( |_, _, _ | Err ( std:: io:: Error :: other ( "DB connection failed" ) ) ) ;
230+ . expect_get_first ( )
231+ . returning ( |_, _| Err ( std:: io:: Error :: other ( "DB connection failed" ) ) ) ;
295232
296233 let result =
297234 get_block_to_restream ( & mock_repo, height, round, valid_round, BlockHash :: default ( ) )
0 commit comments