github.com/muhammedhassanm/blockchain@v0.0.0-20200120143007-697261defd4d/sawtooth-core-master/perf/sawtooth_perf/src/batch_map.rs (about)

     1  /*
     2   * Copyright 2018 Intel Corporation
     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   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   * ------------------------------------------------------------------------------
    16   */
    17  
    18  use std::collections::HashMap;
    19  
    20  use sawtooth_sdk::messages::batch::BatchList;
    21  
    22  pub struct BatchMap {
    23      // Batches that haven't been successfully submitted to
    24      // a validator.
    25      batches_by_id: HashMap<String, BatchList>,
    26  }
    27  
    28  impl BatchMap {
    29      pub fn new() -> BatchMap {
    30          BatchMap {
    31              batches_by_id: HashMap::new(),
    32          }
    33      }
    34  
    35      // Mark that batchlist associated with batch id has been submitted
    36      // to a validator.
    37      pub fn mark_submit_success(&mut self, batch_id: &str) {
    38          self.batches_by_id.remove(batch_id);
    39      }
    40  
    41      // Get a batchlist by id, to submit it to a validator.
    42      pub fn get_batchlist_to_submit(&mut self, batch_id: &str) -> Option<BatchList> {
    43          self.batches_by_id.get(batch_id).cloned()
    44      }
    45  
    46      // Idempotent method for adding a BatchList
    47      pub fn add(&mut self, batchlist: BatchList) {
    48          batchlist
    49              .batches
    50              .last()
    51              .map(|b| b.header_signature.clone())
    52              .map(|batch_id| {
    53                  if !self.batches_by_id.contains_key(batch_id.as_str()) {
    54                      self.batches_by_id.insert(batch_id, batchlist);
    55                  }
    56              });
    57      }
    58  }
    59  
    60  #[cfg(test)]
    61  mod tests {
    62      use super::BatchMap;
    63  
    64      use rand::Rng;
    65      use rand::StdRng;
    66  
    67      use protobuf::RepeatedField;
    68  
    69      use sawtooth_sdk::messages::batch::Batch;
    70      use sawtooth_sdk::messages::batch::BatchList;
    71      use sawtooth_sdk::signing;
    72  
    73      #[test]
    74      fn test_2_cycles_of_retries() {
    75          let mut timed_batch_iterator = BatchMap::new();
    76          let mut batchlists = generate_batchlists(3);
    77  
    78          let batchlist1 = batchlists.pop();
    79          let batchlist2 = batchlists.pop();
    80          let batchlist3 = batchlists.pop();
    81  
    82          let batch_id1 = batchlist1
    83              .clone()
    84              .unwrap()
    85              .batches
    86              .last()
    87              .unwrap()
    88              .header_signature
    89              .clone();
    90          let batch_id2 = batchlist2
    91              .clone()
    92              .unwrap()
    93              .batches
    94              .last()
    95              .unwrap()
    96              .header_signature
    97              .clone();
    98          let batch_id3 = batchlist3
    99              .clone()
   100              .unwrap()
   101              .batches
   102              .last()
   103              .unwrap()
   104              .header_signature
   105              .clone();
   106  
   107          timed_batch_iterator.add(batchlist1.clone().unwrap());
   108          timed_batch_iterator.add(batchlist2.clone().unwrap());
   109          timed_batch_iterator.add(batchlist3.clone().unwrap());
   110  
   111          timed_batch_iterator.add(batchlist1.clone().unwrap());
   112          timed_batch_iterator.add(batchlist2.clone().unwrap());
   113          timed_batch_iterator.add(batchlist3.clone().unwrap());
   114  
   115          timed_batch_iterator.mark_submit_success(&batch_id1);
   116          timed_batch_iterator.mark_submit_success(&batch_id3);
   117  
   118          assert_eq!(
   119              timed_batch_iterator.get_batchlist_to_submit(&batch_id2),
   120              batchlist2
   121          );
   122          assert_eq!(
   123              timed_batch_iterator.get_batchlist_to_submit(&batch_id1),
   124              None
   125          );
   126          assert_eq!(
   127              timed_batch_iterator.get_batchlist_to_submit(&batch_id3),
   128              None
   129          );
   130  
   131          timed_batch_iterator.mark_submit_success(&batch_id2);
   132  
   133          assert_eq!(
   134              timed_batch_iterator.get_batchlist_to_submit(&batch_id2),
   135              None
   136          );
   137      }
   138  
   139      fn generate_batchlists(num: u32) -> Vec<BatchList> {
   140          let context = signing::create_context("secp256k1").unwrap();
   141          let private_key = context.new_random_private_key().unwrap();
   142          let _signer = signing::Signer::new(context.as_ref(), private_key.as_ref());
   143  
   144          let mut batchlists = Vec::new();
   145  
   146          let mut rng = StdRng::new().unwrap();
   147  
   148          for _ in 0..num {
   149              let mut batch = Batch::new();
   150              let mut batchlist = BatchList::new();
   151  
   152              batch.set_header_signature(rng.gen_iter::<char>().take(100).collect());
   153  
   154              let batches = RepeatedField::from_vec(vec![batch]);
   155  
   156              batchlist.set_batches(batches);
   157  
   158              batchlists.push(batchlist);
   159          }
   160          batchlists
   161      }
   162  }