pravega_controller_client/
mock_controller.rs

1/*
2 * Copyright (c) Dell Inc., or its subsidiaries. All Rights Reserved.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 */
10#![allow(dead_code)]
11#![allow(bare_trait_objects)]
12
13use super::ControllerClient;
14use super::ControllerError;
15use crate::ResultRetry;
16use async_trait::async_trait;
17use im::HashMap as ImHashMap;
18use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
19use ordered_float::OrderedFloat;
20use pravega_client_config::connection_type::{ConnectionType, MockType};
21use pravega_client_retry::retry_result::RetryError;
22use pravega_client_shared::*;
23use pravega_connection_pool::connection_pool::ConnectionPool;
24use pravega_wire_protocol::client_connection::{ClientConnection, ClientConnectionImpl};
25use pravega_wire_protocol::commands::{CreateSegmentCommand, DeleteSegmentCommand, MergeSegmentsCommand};
26use pravega_wire_protocol::connection_factory::{
27    ConnectionFactory, ConnectionFactoryConfig, SegmentConnectionManager,
28};
29use pravega_wire_protocol::error::ClientConnectionError;
30use pravega_wire_protocol::wire_commands::{Replies, Requests};
31use serde::{Deserialize, Serialize};
32use std::collections::HashSet;
33use std::collections::{BTreeMap, HashMap};
34use std::sync::atomic::{AtomicUsize, Ordering};
35use std::time::Duration;
36use std::time::SystemTime;
37use tokio::sync::{RwLock, RwLockReadGuard};
38use uuid::Uuid;
39
40static ID_GENERATOR: AtomicUsize = AtomicUsize::new(0);
41
42pub struct MockController {
43    endpoint: PravegaNodeUri,
44    pool: ConnectionPool<SegmentConnectionManager>,
45    created_scopes: RwLock<HashMap<String, HashSet<ScopedStream>>>,
46    created_streams: RwLock<HashMap<ScopedStream, StreamConfiguration>>,
47    transactions: RwLock<HashMap<TxId, TransactionStatus>>,
48}
49
50impl MockController {
51    pub fn new(endpoint: PravegaNodeUri) -> Self {
52        let config = ConnectionFactoryConfig::new(ConnectionType::Mock(MockType::Happy));
53        let cf = ConnectionFactory::create(config) as Box<dyn ConnectionFactory>;
54        let manager = SegmentConnectionManager::new(cf, 10);
55        let pool = ConnectionPool::new(manager);
56        MockController {
57            endpoint,
58            pool,
59            created_scopes: RwLock::new(HashMap::new()),
60            created_streams: RwLock::new(HashMap::new()),
61            transactions: RwLock::new(HashMap::new()),
62        }
63    }
64}
65#[async_trait]
66impl ControllerClient for MockController {
67    async fn create_scope(&self, scope: &Scope) -> Result<bool, RetryError<ControllerError>> {
68        let scope_name = scope.name.clone();
69        if self.created_scopes.read().await.contains_key(&scope_name) {
70            return Ok(false);
71        }
72
73        self.created_scopes
74            .write()
75            .await
76            .insert(scope_name, HashSet::new());
77        Ok(true)
78    }
79
80    async fn check_scope_exists(&self, scope: &Scope) -> Result<bool, RetryError<ControllerError>> {
81        let scope_name = scope.name.clone();
82        if self.created_scopes.read().await.contains_key(&scope_name) {
83            return Ok(true);
84        }
85        Ok(false)
86    }
87
88    async fn list_scopes(
89        &self,
90        _token: &CToken,
91    ) -> Result<Option<(Vec<Scope>, CToken)>, RetryError<ControllerError>> {
92        let map_guard = self.created_scopes.read().await;
93        let scopes = map_guard.keys();
94        let mut result = Vec::new();
95        for scope in scopes {
96            result.push(Scope::from(scope.clone()));
97        }
98        Ok(Some((result, CToken::from("mock_token"))))
99    }
100
101    async fn list_streams(
102        &self,
103        scope: &Scope,
104        _token: &CToken,
105    ) -> Result<Option<(Vec<ScopedStream>, CToken)>, RetryError<ControllerError>> {
106        let map_guard = self.created_scopes.read().await;
107        let streams_set = map_guard.get(&scope.name).ok_or(RetryError {
108            error: ControllerError::OperationError {
109                can_retry: false,
110                operation: "listStreams".into(),
111                error_msg: "Scope not exist".into(),
112            },
113            total_delay: Duration::from_millis(1),
114            tries: 0,
115        })?;
116        let mut result = Vec::new();
117        for stream in streams_set {
118            result.push(stream.clone())
119        }
120        Ok(Some((result, CToken::from("mock_token"))))
121    }
122
123    async fn list_streams_for_tag(
124        &self,
125        scope: &Scope,
126        tag: &str,
127        _token: &CToken,
128    ) -> Result<Option<(Vec<ScopedStream>, CToken)>, RetryError<ControllerError>> {
129        let scope_gaurd = self.created_scopes.read().await;
130        let stream_gaurd = self.created_streams.read().await;
131
132        let streams_set = scope_gaurd.get(&scope.name).ok_or(RetryError {
133            error: ControllerError::OperationError {
134                can_retry: false,
135                operation: "listStreams".into(),
136                error_msg: "Scope not exist".into(),
137            },
138            total_delay: Duration::from_millis(1),
139            tries: 0,
140        })?;
141        let mut result = Vec::new();
142        for stream in streams_set {
143            let cfg = stream_gaurd.get(stream).ok_or(RetryError {
144                error: ControllerError::OperationError {
145                    can_retry: false,
146                    operation: "listStreamsForTag".into(),
147                    error_msg: "Stream does not exist".into(),
148                },
149                total_delay: Duration::from_millis(1),
150                tries: 0,
151            })?;
152            match &cfg.tags {
153                None => {}
154                Some(tag_list) => {
155                    if tag_list.contains(&tag.to_string()) {
156                        result.push(stream.clone())
157                    }
158                }
159            };
160        }
161        Ok(Some((result, CToken::from("mock_token"))))
162    }
163
164    async fn delete_scope(&self, scope: &Scope) -> Result<bool, RetryError<ControllerError>> {
165        let scope_name = scope.name.clone();
166        if self.created_scopes.read().await.get(&scope_name).is_none() {
167            return Ok(false);
168        }
169
170        if !self
171            .created_scopes
172            .read()
173            .await
174            .get(&scope_name)
175            .unwrap()
176            .is_empty()
177        {
178            Err(RetryError {
179                error: ControllerError::OperationError {
180                    can_retry: false,
181                    operation: "DeleteScope".into(),
182                    error_msg: "Scope not empty".into(),
183                },
184                total_delay: Duration::from_millis(1),
185                tries: 0,
186            })
187        } else {
188            self.created_scopes.write().await.remove(&scope_name);
189            Ok(true)
190        }
191    }
192
193    async fn create_stream(
194        &self,
195        stream_config: &StreamConfiguration,
196    ) -> Result<bool, RetryError<ControllerError>> {
197        let stream = stream_config.scoped_stream.clone();
198        if self.created_streams.read().await.contains_key(&stream) {
199            return Ok(false);
200        }
201        if self.created_scopes.read().await.get(&stream.scope.name).is_none() {
202            return Err(RetryError {
203                error: ControllerError::OperationError {
204                    can_retry: false,
205                    operation: "create stream".into(),
206                    error_msg: "Scope does not exist.".into(),
207                },
208                total_delay: Duration::from_millis(1),
209                tries: 0,
210            });
211        }
212        self.created_streams
213            .write()
214            .await
215            .insert(stream.clone(), stream_config.clone());
216        self.created_scopes
217            .write()
218            .await
219            .get_mut(&stream.scope.name)
220            .unwrap()
221            .insert(stream.clone());
222
223        let read_guard = &self.created_streams.read().await;
224        for segment in get_segments_for_stream(&stream, read_guard)? {
225            let segment_name = segment.to_string();
226            create_segment(segment_name, self, false).await?;
227        }
228        Ok(true)
229    }
230
231    async fn check_stream_exists(&self, stream: &ScopedStream) -> Result<bool, RetryError<ControllerError>> {
232        if self.created_streams.read().await.contains_key(stream) {
233            return Ok(true);
234        }
235        Ok(false)
236    }
237
238    async fn update_stream(
239        &self,
240        _stream_config: &StreamConfiguration,
241    ) -> Result<bool, RetryError<ControllerError>> {
242        Err(RetryError {
243            error: ControllerError::OperationError {
244                can_retry: false,
245                operation: "update stream".into(),
246                error_msg: "unsupported operation.".into(),
247            },
248            total_delay: Duration::from_millis(1),
249            tries: 0,
250        })
251    }
252
253    async fn get_stream_configuration(&self, _stream: &ScopedStream) -> ResultRetry<StreamConfiguration> {
254        Err(RetryError {
255            error: ControllerError::OperationError {
256                can_retry: false,
257                operation: "get stream configuration".into(),
258                error_msg: "unsupported operation.".into(),
259            },
260            total_delay: Duration::from_millis(1),
261            tries: 0,
262        })
263    }
264
265    async fn get_stream_tags(&self, _stream: &ScopedStream) -> ResultRetry<Option<Vec<String>>> {
266        Err(RetryError {
267            error: ControllerError::OperationError {
268                can_retry: false,
269                operation: "get stream tags".into(),
270                error_msg: "unsupported operation.".into(),
271            },
272            total_delay: Duration::from_millis(1),
273            tries: 0,
274        })
275    }
276
277    async fn truncate_stream(&self, _stream_cut: &StreamCut) -> Result<bool, RetryError<ControllerError>> {
278        Err(RetryError {
279            error: ControllerError::OperationError {
280                can_retry: false,
281                operation: "truncate stream".into(),
282                error_msg: "unsupported operation.".into(),
283            },
284            total_delay: Duration::from_millis(1),
285            tries: 0,
286        })
287    }
288
289    async fn seal_stream(&self, _stream: &ScopedStream) -> Result<bool, RetryError<ControllerError>> {
290        Err(RetryError {
291            error: ControllerError::OperationError {
292                can_retry: false,
293                operation: "seal stream".into(),
294                error_msg: "unsupported operation.".into(),
295            },
296            total_delay: Duration::from_millis(1),
297            tries: 0,
298        })
299    }
300
301    async fn delete_stream(&self, stream: &ScopedStream) -> Result<bool, RetryError<ControllerError>> {
302        if self.created_streams.read().await.get(stream).is_none() {
303            return Ok(false);
304        }
305
306        for segment in get_segments_for_stream(stream, &self.created_streams.read().await)? {
307            let segment_name = segment.to_string();
308            delete_segment(segment_name, self, false).await?;
309        }
310
311        self.created_streams.write().await.remove(stream);
312        self.created_scopes
313            .write()
314            .await
315            .get_mut(&stream.scope.name)
316            .unwrap()
317            .remove(stream);
318        Ok(true)
319    }
320
321    async fn get_current_segments(
322        &self,
323        stream: &ScopedStream,
324    ) -> Result<StreamSegments, RetryError<ControllerError>> {
325        let segments_in_stream = get_segments_for_stream(stream, &self.created_streams.read().await)?;
326        let mut segments = BTreeMap::new();
327        let increment = 1.0 / segments_in_stream.len() as f64;
328        for (number, segment) in segments_in_stream.into_iter().enumerate() {
329            let segment_with_range = SegmentWithRange {
330                scoped_segment: segment,
331                min_key: OrderedFloat(number as f64 * increment),
332                max_key: OrderedFloat((number + 1) as f64 * increment),
333            };
334            segments.insert(segment_with_range.max_key, segment_with_range);
335        }
336
337        Ok(StreamSegments {
338            key_segment_map: segments.into(),
339        })
340    }
341
342    async fn get_epoch_segments(
343        &self,
344        stream: &ScopedStream,
345        _epoch: i32,
346    ) -> Result<StreamSegments, RetryError<ControllerError>> {
347        let segments_in_stream = get_segments_for_stream(stream, &self.created_streams.read().await)?;
348        let mut segments = BTreeMap::new();
349        let increment = 1.0 / segments_in_stream.len() as f64;
350        for (number, segment) in segments_in_stream.into_iter().enumerate() {
351            let segment_with_range = SegmentWithRange {
352                scoped_segment: segment,
353                min_key: OrderedFloat(number as f64 * increment),
354                max_key: OrderedFloat((number + 1) as f64 * increment),
355            };
356            segments.insert(segment_with_range.max_key, segment_with_range);
357        }
358
359        Ok(StreamSegments {
360            key_segment_map: segments.into(),
361        })
362    }
363
364    async fn get_head_segments(
365        &self,
366        stream: &ScopedStream,
367    ) -> ResultRetry<std::collections::HashMap<Segment, i64>> {
368        let segments_in_stream: Vec<ScopedSegment> =
369            get_segments_for_stream(stream, &self.created_streams.read().await)?;
370        Ok(segments_in_stream
371            .iter()
372            .map(|t| (t.segment.clone(), 0i64))
373            .collect())
374    }
375
376    async fn create_transaction(
377        &self,
378        stream: &ScopedStream,
379        _lease: Duration,
380    ) -> Result<TxnSegments, RetryError<ControllerError>> {
381        let uuid = Uuid::new_v4().as_u128();
382        let current_segments = self.get_current_segments(stream).await?;
383        let mut guard = self.transactions.write().await;
384        guard.insert(TxId(uuid), TransactionStatus::Open);
385
386        Ok(TxnSegments {
387            stream_segments: current_segments,
388            tx_id: TxId(uuid),
389        })
390    }
391
392    async fn ping_transaction(
393        &self,
394        _stream: &ScopedStream,
395        tx_id: TxId,
396        _lease: Duration,
397    ) -> Result<PingStatus, RetryError<ControllerError>> {
398        let guard = self.transactions.read().await;
399        let status = guard.get(&tx_id).expect("get transaction status");
400        match status {
401            TransactionStatus::Committed => Ok(PingStatus::Committed),
402            TransactionStatus::Aborted => Ok(PingStatus::Aborted),
403            _ => Ok(PingStatus::Ok),
404        }
405    }
406
407    async fn commit_transaction(
408        &self,
409        _stream: &ScopedStream,
410        tx_id: TxId,
411        _writer_id: WriterId,
412        _time: Timestamp,
413    ) -> Result<(), RetryError<ControllerError>> {
414        let mut guard = self.transactions.write().await;
415        guard.insert(tx_id, TransactionStatus::Committed);
416        Ok(())
417    }
418
419    async fn abort_transaction(
420        &self,
421        _stream: &ScopedStream,
422        tx_id: TxId,
423    ) -> Result<(), RetryError<ControllerError>> {
424        let mut guard = self.transactions.write().await;
425        guard.insert(tx_id, TransactionStatus::Aborted);
426        Ok(())
427    }
428
429    async fn check_transaction_status(
430        &self,
431        _stream: &ScopedStream,
432        tx_id: TxId,
433    ) -> Result<TransactionStatus, RetryError<ControllerError>> {
434        let guard = self.transactions.read().await;
435        let status = guard.get(&tx_id).expect("get transaction");
436        Ok(status.clone())
437    }
438
439    async fn get_endpoint_for_segment(
440        &self,
441        _segment: &ScopedSegment,
442    ) -> Result<PravegaNodeUri, RetryError<ControllerError>> {
443        Ok(self.endpoint.clone())
444    }
445
446    async fn get_or_refresh_delegation_token_for(
447        &self,
448        _stream: ScopedStream,
449    ) -> Result<String, RetryError<ControllerError>> {
450        let now = SystemTime::now()
451            .duration_since(SystemTime::UNIX_EPOCH)
452            .expect("get unix time");
453        let timeout = Duration::from_secs(5);
454        let expiry_time = now.checked_add(timeout).expect("calculate expiry time");
455        let claims = Claims {
456            sub: "subject".to_string(),
457            aud: "segmentstore".to_string(),
458            iat: now.as_secs(),
459            exp: expiry_time.as_secs(),
460        };
461
462        let header = Header {
463            typ: Some("JWT".to_owned()),
464            alg: Algorithm::HS256,
465            ..Default::default()
466        };
467
468        let key = b"secret";
469        let token = encode(&header, &claims, &EncodingKey::from_secret(key)).expect("encode to JWT token");
470        Ok(token)
471    }
472
473    async fn get_successors(
474        &self,
475        _segment: &ScopedSegment,
476    ) -> Result<StreamSegmentsWithPredecessors, RetryError<ControllerError>> {
477        // empty hash map means the stream is sealed
478        Ok(StreamSegmentsWithPredecessors {
479            segment_with_predecessors: ImHashMap::new(),
480            replacement_segments: ImHashMap::new(),
481        })
482    }
483
484    async fn scale_stream(
485        &self,
486        _stream: &ScopedStream,
487        _sealed_segments: &[Segment],
488        _new_key_ranges: &[(f64, f64)],
489    ) -> Result<(), RetryError<ControllerError>> {
490        Err(RetryError {
491            error: ControllerError::OperationError {
492                can_retry: false, // do not retry.
493                operation: "scale stream".into(),
494                error_msg: "unsupported operation.".into(),
495            },
496            total_delay: Duration::from_millis(1),
497            tries: 0,
498        })
499    }
500
501    async fn check_scale(
502        &self,
503        _stream: &ScopedStream,
504        _scale_epoch: i32,
505    ) -> Result<bool, RetryError<ControllerError>> {
506        Err(RetryError {
507            error: ControllerError::OperationError {
508                can_retry: false, // do not retry.
509                operation: "check stream scale".into(),
510                error_msg: "unsupported operation.".into(),
511            },
512            total_delay: Duration::from_millis(1),
513            tries: 0,
514        })
515    }
516}
517
518fn get_segments_for_stream(
519    stream: &ScopedStream,
520    created_streams: &RwLockReadGuard<HashMap<ScopedStream, StreamConfiguration>>,
521) -> Result<Vec<ScopedSegment>, RetryError<ControllerError>> {
522    let stream_config = created_streams.get(stream);
523    if stream_config.is_none() {
524        return Err(RetryError {
525            error: ControllerError::OperationError {
526                can_retry: false, // do not retry.
527                operation: "get segments for stream".into(),
528                error_msg: "stream does not exist.".into(),
529            },
530            total_delay: Duration::from_millis(1),
531            tries: 0,
532        });
533    }
534
535    let scaling_policy = stream_config.unwrap().scaling.clone();
536
537    if scaling_policy.scale_type != ScaleType::FixedNumSegments {
538        return Err(RetryError {
539            error: ControllerError::OperationError {
540                can_retry: false, // do not retry.
541                operation: "get segments for stream".into(),
542                error_msg: "Dynamic scaling not supported with a mock controller.".into(),
543            },
544            total_delay: Duration::from_millis(1),
545            tries: 0,
546        });
547    }
548    let mut result = Vec::with_capacity(scaling_policy.min_num_segments as usize);
549    for i in 0..scaling_policy.min_num_segments {
550        result.push(ScopedSegment {
551            scope: stream.scope.clone(),
552            stream: stream.stream.clone(),
553            segment: Segment::from(i.into()),
554        });
555    }
556
557    Ok(result)
558}
559
560async fn create_segment(
561    name: String,
562    controller: &MockController,
563    call_server: bool,
564) -> Result<bool, RetryError<ControllerError>> {
565    if !call_server {
566        return Ok(true);
567    }
568    let scale_type = ScaleType::FixedNumSegments;
569    let id = ID_GENERATOR.fetch_add(1, Ordering::SeqCst) as i64;
570    let command = Requests::CreateSegment(CreateSegmentCommand {
571        request_id: id,
572        segment: name,
573        target_rate: 0,
574        scale_type: scale_type as u8,
575        delegation_token: String::from(""),
576    });
577    let reply = send_request_over_connection(&command, controller).await;
578    match reply {
579        Ok(r) => {
580            match r {
581                Replies::WrongHost(_) => Err(ControllerError::OperationError {
582                    can_retry: false, // do not retry.
583                    operation: "create segment".into(),
584                    error_msg: "Wrong host.".into(),
585                }),
586                Replies::SegmentCreated(_) => Ok(true),
587                Replies::SegmentAlreadyExists(_) => Ok(false),
588                Replies::AuthTokenCheckFailed(_) => Err(ControllerError::OperationError {
589                    can_retry: false, // do not retry.
590                    operation: "create segment".into(),
591                    error_msg: "authToken check failed,".into(),
592                }),
593                _ => Err(ControllerError::OperationError {
594                    can_retry: false, // do not retry.
595                    operation: "create segment".into(),
596                    error_msg: "Unsupported Command".into(),
597                }),
598            }
599        }
600        Err(_e) => Err(ControllerError::OperationError {
601            can_retry: false, // do not retry.
602            operation: "create segment".into(),
603            error_msg: "Connection Error".into(),
604        }),
605    }
606    .map_err(|e| RetryError {
607        error: e,
608        total_delay: Duration::from_millis(1),
609        tries: 0,
610    })
611}
612
613async fn delete_segment(
614    name: String,
615    controller: &MockController,
616    call_server: bool,
617) -> Result<bool, RetryError<ControllerError>> {
618    if !call_server {
619        return Ok(true);
620    }
621    let id = ID_GENERATOR.fetch_add(1, Ordering::SeqCst) as i64;
622    let command = Requests::DeleteSegment(DeleteSegmentCommand {
623        request_id: id,
624        segment: name,
625        delegation_token: String::from(""),
626    });
627    let reply = send_request_over_connection(&command, controller).await;
628    match reply {
629        Err(_e) => Err(ControllerError::OperationError {
630            can_retry: false, // do not retry.
631            operation: "delete segment".into(),
632            error_msg: "Connection Error".into(),
633        }),
634        Ok(r) => {
635            match r {
636                Replies::WrongHost(_) => Err(ControllerError::OperationError {
637                    can_retry: false, // do not retry.
638                    operation: "delete segment".into(),
639                    error_msg: "Wrong host.".into(),
640                }),
641                Replies::SegmentDeleted(_) => Ok(true),
642                Replies::NoSuchSegment(_) => Ok(false),
643                Replies::AuthTokenCheckFailed(_) => Err(ControllerError::OperationError {
644                    can_retry: false, // do not retry.
645                    operation: "delete segment".into(),
646                    error_msg: "authToken check failed,".into(),
647                }),
648                _ => Err(ControllerError::OperationError {
649                    can_retry: false, // do not retry.
650                    operation: "delete segment".into(),
651                    error_msg: "Unsupported Command.".into(),
652                }),
653            }
654        }
655    }
656    .map_err(|e| RetryError {
657        error: e,
658        total_delay: Duration::from_millis(1),
659        tries: 0,
660    })
661}
662
663async fn commit_tx_segment(
664    uuid: TxId,
665    segment: ScopedSegment,
666    controller: &MockController,
667    call_server: bool,
668) -> Result<(), RetryError<ControllerError>> {
669    if !call_server {
670        return Ok(());
671    }
672    let source_name = segment.scope.name.clone() + &uuid.to_string();
673    let id = ID_GENERATOR.fetch_add(1, Ordering::SeqCst) as i64;
674    let command = Requests::MergeSegments(MergeSegmentsCommand {
675        request_id: id,
676        target: segment.to_string(),
677        source: source_name,
678        delegation_token: String::from(""),
679    });
680    let reply = send_request_over_connection(&command, controller).await;
681    match reply {
682        Err(_e) => Err(ControllerError::OperationError {
683            can_retry: false, // do not retry.
684            operation: "commit tx segment".into(),
685            error_msg: "Connection Error".into(),
686        }),
687        Ok(r) => {
688            match r {
689                Replies::SegmentsMerged(_) => Ok(()),
690                Replies::WrongHost(_) => Err(ControllerError::OperationError {
691                    can_retry: false, // do not retry.
692                    operation: "commit tx segment".into(),
693                    error_msg: "Wrong host.".into(),
694                }),
695                Replies::SegmentDeleted(_) => Err(ControllerError::OperationError {
696                    can_retry: false, // do not retry.
697                    operation: "commit tx segment".into(),
698                    error_msg: "Transaction already aborted.".into(),
699                }),
700                Replies::AuthTokenCheckFailed(_) => Err(ControllerError::OperationError {
701                    can_retry: false, // do not retry.
702                    operation: "commit tx segment".into(),
703                    error_msg: "authToken check failed,".into(),
704                }),
705                _ => Err(ControllerError::OperationError {
706                    can_retry: false, // do not retry.
707                    operation: "commit tx segment".into(),
708                    error_msg: "Unsupported Command,".into(),
709                }),
710            }
711        }
712    }
713    .map_err(|e| RetryError {
714        error: e,
715        total_delay: Duration::from_millis(1),
716        tries: 0,
717    })
718}
719
720async fn abort_tx_segment(
721    uuid: TxId,
722    segment: ScopedSegment,
723    controller: &MockController,
724    call_server: bool,
725) -> Result<(), RetryError<ControllerError>> {
726    if !call_server {
727        return Ok(());
728    }
729    let transaction_name = segment.scope.name.clone() + &uuid.to_string();
730    let id = ID_GENERATOR.fetch_add(1, Ordering::SeqCst) as i64;
731    let command = Requests::DeleteSegment(DeleteSegmentCommand {
732        request_id: id,
733        segment: transaction_name,
734        delegation_token: String::from(""),
735    });
736    let reply = send_request_over_connection(&command, controller).await;
737    match reply {
738        Err(_e) => Err(ControllerError::OperationError {
739            can_retry: false, // do not retry.
740            operation: "abort tx segment".into(),
741            error_msg: "Connection Error".into(),
742        }),
743        Ok(r) => {
744            match r {
745                Replies::SegmentsMerged(_) => Err(ControllerError::OperationError {
746                    can_retry: false, // do not retry.
747                    operation: "abort tx segment".into(),
748                    error_msg: "Transaction already committed.".into(),
749                }),
750                Replies::WrongHost(_) => Err(ControllerError::OperationError {
751                    can_retry: false, // do not retry.
752                    operation: "abort tx segment".into(),
753                    error_msg: "Wrong host.".into(),
754                }),
755                Replies::AuthTokenCheckFailed(_) => Err(ControllerError::OperationError {
756                    can_retry: false, // do not retry.
757                    operation: "abort tx segment".into(),
758                    error_msg: "authToken check failed,".into(),
759                }),
760                Replies::SegmentDeleted(_) => Ok(()),
761                _ => Err(ControllerError::OperationError {
762                    can_retry: false, // do not retry.
763                    operation: "abort tx segment".into(),
764                    error_msg: "Unsupported Command,".into(),
765                }),
766            }
767        }
768    }
769    .map_err(|e| RetryError {
770        error: e,
771        total_delay: Duration::from_millis(1),
772        tries: 0,
773    })
774}
775
776async fn create_tx_segment(
777    uuid: TxId,
778    segment: ScopedSegment,
779    controller: &MockController,
780    call_server: bool,
781) -> Result<(), RetryError<ControllerError>> {
782    if !call_server {
783        return Ok(());
784    }
785    let transaction_name = segment.scope.name.clone() + &uuid.to_string();
786    let scale_type = ScaleType::FixedNumSegments;
787    let id = ID_GENERATOR.fetch_add(1, Ordering::SeqCst) as i64;
788    let command = Requests::CreateSegment(CreateSegmentCommand {
789        request_id: id,
790        segment: transaction_name,
791        target_rate: 0,
792        scale_type: scale_type as u8,
793        delegation_token: String::from(""),
794    });
795    let reply = send_request_over_connection(&command, controller).await;
796    match reply {
797        Err(_e) => Err(ControllerError::OperationError {
798            can_retry: false, // do not retry.
799            operation: "abort tx segment".into(),
800            error_msg: "Connection Error".into(),
801        }),
802        Ok(r) => {
803            match r {
804                Replies::SegmentCreated(_) => Ok(()),
805                Replies::WrongHost(_) => Err(ControllerError::OperationError {
806                    can_retry: false, // do not retry.
807                    operation: "create tx segment".into(),
808                    error_msg: "Wrong host.".into(),
809                }),
810                Replies::AuthTokenCheckFailed(_) => Err(ControllerError::OperationError {
811                    can_retry: false, // do not retry.
812                    operation: "create tx segment".into(),
813                    error_msg: "authToken check failed,".into(),
814                }),
815                _ => Err(ControllerError::OperationError {
816                    can_retry: false, // do not retry.
817                    operation: "create tx segment".into(),
818                    error_msg: "Unsupported Command,".into(),
819                }),
820            }
821        }
822    }
823    .map_err(|e| RetryError {
824        error: e,
825        total_delay: Duration::from_millis(1),
826        tries: 0,
827    })
828}
829
830async fn send_request_over_connection(
831    command: &Requests,
832    controller: &MockController,
833) -> Result<Replies, ClientConnectionError> {
834    let pooled_connection = controller
835        .pool
836        .get_connection(controller.endpoint.clone())
837        .await
838        .expect("get connection from pool");
839    let mut connection = ClientConnectionImpl {
840        connection: pooled_connection,
841    };
842    connection.write(command).await?;
843    connection.read().await
844}
845
846#[derive(Debug, Serialize, Deserialize)]
847struct Claims {
848    sub: String,
849    aud: String,
850    iat: u64,
851    exp: u64,
852}
pravega_controller_client/mock_controller.rs

pravega_controller_client/
mock_controller.rs
