From 6f4a98f6fd5c7b484f8397b993b4a4b463ccaeb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire=20Henry?= <gregoire.henry@tezos.com>
Date: Sat, 2 Jun 2018 18:30:40 +0200
Subject: [PATCH] P2P: allow to limit the size of mempool

---
 src/lib_base/mempool.ml                  | 10 +++++++++-
 src/lib_base/mempool.mli                 |  1 +
 src/lib_shell/distributed_db_message.ml  | 12 +++++++++++-
 src/lib_shell/distributed_db_message.mli |  1 +
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/lib_base/mempool.ml b/src/lib_base/mempool.ml
index 3d2a4492b..2688b3f11 100644
--- a/src/lib_base/mempool.ml
+++ b/src/lib_base/mempool.ml
@@ -19,9 +19,17 @@ let encoding =
     (fun { known_valid ; pending } -> (known_valid, pending))
     (fun (known_valid, pending) -> { known_valid ; pending })
     (obj2
-       (req "known_valid" (dynamic_size (list Operation_hash.encoding)))
+       (req "known_valid" (list Operation_hash.encoding))
        (req "pending" (dynamic_size Operation_hash.Set.encoding)))
 
+let bounded_encoding ?max_operations () =
+  match max_operations with
+  | None -> encoding
+  | Some max_operations ->
+      Data_encoding.check_size
+        (8 + max_operations * Operation_hash.size)
+        encoding
+
 let empty = {
   known_valid = [] ;
   pending = Operation_hash.Set.empty ;
diff --git a/src/lib_base/mempool.mli b/src/lib_base/mempool.mli
index 0c13ccf91..40a157e43 100644
--- a/src/lib_base/mempool.mli
+++ b/src/lib_base/mempool.mli
@@ -19,6 +19,7 @@ type t = {
 type mempool = t
 
 val encoding: mempool Data_encoding.t
+val bounded_encoding: ?max_operations:int -> unit -> mempool Data_encoding.t
 
 val empty: mempool
 (** Empty mempool. *)
diff --git a/src/lib_shell/distributed_db_message.ml b/src/lib_shell/distributed_db_message.ml
index d3506357f..bb955fa50 100644
--- a/src/lib_shell/distributed_db_message.ml
+++ b/src/lib_shell/distributed_db_message.ml
@@ -93,6 +93,16 @@ module Bounded_encoding = struct
     protocol_max_size := max
   let protocol = delayed (fun () -> !protocol_cache)
 
+  let mempool_max_operations = ref None
+  let mempool_cache =
+    ref (Mempool.bounded_encoding ?max_operations:!mempool_max_operations ())
+  let update_mempool_encoding () =
+    mempool_cache :=
+      Mempool.bounded_encoding ?max_operations:!mempool_max_operations ()
+  let set_mempool_max_operations max =
+    mempool_max_operations := max
+  let mempool = delayed (fun () -> !mempool_cache)
+
 end
 
 type t =
@@ -170,7 +180,7 @@ let encoding =
       (obj3
          (req "chain_id" Chain_id.encoding)
          (req "current_block_header" (dynamic_size Bounded_encoding.block_header))
-         (req "current_mempool" Mempool.encoding))
+         (req "current_mempool" Bounded_encoding.mempool))
       (function
         | Current_head (chain_id, bh, mempool) -> Some (chain_id, bh, mempool)
         | _ -> None)
diff --git a/src/lib_shell/distributed_db_message.mli b/src/lib_shell/distributed_db_message.mli
index 3ce24c929..32a932af4 100644
--- a/src/lib_shell/distributed_db_message.mli
+++ b/src/lib_shell/distributed_db_message.mli
@@ -48,4 +48,5 @@ module Bounded_encoding : sig
   val set_operation_list_max_length: int option -> unit
   val set_operation_max_pass: int option -> unit
   val set_protocol_max_size: int option -> unit
+  val set_mempool_max_operations: int option -> unit
 end