4 files changed, 56 insertions, 11 deletions

diff --git a/src/vnet/ip/ip4.h b/src/vnet/ip/ip4.h
index 4c5cc054c4a..f7e5fe332b1 100644
--- a/src/vnet/ip/ip4.h
+++ b/src/vnet/ip/ip4.h
@@ -148,6 +148,12 @@ typedef struct ip4_main_t
 
     u8 pad[2];
   } host_config;
+
+  /** Heapsize for the Mtries */
+  uword mtrie_heap_size;
+
+  /** The memory heap for the mtries */
+  void *mtrie_mheap;
 } ip4_main_t;
 
 /** Global ip4 main structure. */
diff --git a/src/vnet/ip/ip4_forward.c b/src/vnet/ip/ip4_forward.c
index d7afc84ea05..c7cf362a3c0 100755
--- a/src/vnet/ip/ip4_forward.c
+++ b/src/vnet/ip/ip4_forward.c
@@ -1189,6 +1189,12 @@ ip4_lookup_init (vlib_main_t * vm)
 
   if ((error = vlib_call_init_function (vm, vnet_feature_init)))
     return error;
+  if ((error = vlib_call_init_function (vm, ip4_mtrie_module_init)))
+    return (error);
+  if ((error = vlib_call_init_function (vm, fib_module_init)))
+    return error;
+  if ((error = vlib_call_init_function (vm, mfib_module_init)))
+    return error;
 
   for (i = 0; i < ARRAY_LEN (im->fib_masks); i++)
     {
@@ -3166,6 +3172,29 @@ VLIB_CLI_COMMAND (set_ip_classify_command, static) =
 };
 /* *INDENT-ON* */
 
+static clib_error_t *
+ip4_config (vlib_main_t * vm, unformat_input_t * input)
+{
+  ip4_main_t *im = &ip4_main;
+  uword heapsize = 0;
+
+  while (unformat_check_input (input) != UNFORMAT_END_OF_INPUT)
+    {
+      if (unformat (input, "heap-size %U", unformat_memory_size, &heapsize))
+	;
+      else
+	return clib_error_return (0,
+				  "invalid heap-size parameter `%U'",
+				  format_unformat_error, input);
+    }
+
+  im->mtrie_heap_size = heapsize;
+
+  return 0;
+}
+
+VLIB_EARLY_CONFIG_FUNCTION (ip4_config, "ip");
+
 /*
  * fd.io coding-style-patch-verification: ON
  *
diff --git a/src/vnet/ip/ip4_mtrie.c b/src/vnet/ip/ip4_mtrie.c
index 5b78cbec1ea..5981a3c9e20 100644
--- a/src/vnet/ip/ip4_mtrie.c
+++ b/src/vnet/ip/ip4_mtrie.c
@@ -176,9 +176,12 @@ ply_create (ip4_fib_mtrie_t * m,
 	    u32 leaf_prefix_len, u32 ply_base_len)
 {
   ip4_fib_mtrie_8_ply_t *p;
-
+  void *old_heap;
   /* Get cache aligned ply. */
+
+  old_heap = clib_mem_set_heap (ip4_main.mtrie_mheap);
   pool_get_aligned (ip4_ply_pool, p, CLIB_CACHE_LINE_BYTES);
+  clib_mem_set_heap (old_heap);
 
   ply_8_init (p, init_leaf, leaf_prefix_len, ply_base_len);
   return ip4_fib_mtrie_leaf_set_next_ply_index (p - ip4_ply_pool);
@@ -796,15 +799,27 @@ format_ip4_fib_mtrie (u8 * s, va_list * va)
   return s;
 }
 
+/** Default heap size for the IPv4 mtries */
+#define IP4_FIB_DEFAULT_MTRIE_HEAP_SIZE (32<<20)
+
 static clib_error_t *
 ip4_mtrie_module_init (vlib_main_t * vm)
 {
-  /* Burn one ply so index 0 is taken */
   CLIB_UNUSED (ip4_fib_mtrie_8_ply_t * p);
+  ip4_main_t *im = &ip4_main;
+  clib_error_t *error = NULL;
+  uword *old_heap;
+
+  if (0 == im->mtrie_heap_size)
+    im->mtrie_heap_size = IP4_FIB_DEFAULT_MTRIE_HEAP_SIZE;
+  im->mtrie_mheap = mheap_alloc (0, im->mtrie_heap_size);
 
+  /* Burn one ply so index 0 is taken */
+  old_heap = clib_mem_set_heap (ip4_main.mtrie_mheap);
   pool_get (ip4_ply_pool, p);
+  clib_mem_set_heap (old_heap);
 
-  return (NULL);
+  return (error);
 }
 
 VLIB_INIT_FUNCTION (ip4_mtrie_module_init);
diff --git a/src/vnet/ip/ip6_forward.c b/src/vnet/ip/ip6_forward.c
index a14bf09742b..67fedc949ac 100644
--- a/src/vnet/ip/ip6_forward.c
+++ b/src/vnet/ip/ip6_forward.c
@@ -3494,14 +3494,9 @@ ip6_config (vlib_main_t * vm, unformat_input_t * input)
     {
       if (unformat (input, "hash-buckets %d", &tmp))
 	nbuckets = tmp;
-      else if (unformat (input, "heap-size %dm", &tmp))
-	heapsize = ((u64) tmp) << 20;
-      else if (unformat (input, "heap-size %dM", &tmp))
-	heapsize = ((u64) tmp) << 20;
-      else if (unformat (input, "heap-size %dg", &tmp))
-	heapsize = ((u64) tmp) << 30;
-      else if (unformat (input, "heap-size %dG", &tmp))
-	heapsize = ((u64) tmp) << 30;
+      else if (unformat (input, "heap-size %U",
+			 unformat_memory_size, &heapsize))
+	;
       else
 	return clib_error_return (0, "unknown input '%U'",
 				  format_unformat_error, input);