Generate a unique ID for each Netfront request

Previously we used the grant ref, but this isn't unique. This could
possibly cause pages to be returned to the free pool while they were
still being transmitted.

Since IDs must be 16-bit, this patch also limits the size of the pool
(to 256 pages, which is roughly the size of the ring).

Author: talex5

lib/frontend.ml

@@ -283,12 +283,11 @@ module Make(C: S.CONFIGURATION with type 'a io = 'a Lwt.t) = struct
    * remaining (unsent) data and a thread which will return when the data has
    * been ack'd by netback. *)
   let write_request ?size ~flags nf datav =
-    Shared_page_pool.use nf.t.tx_pool (fun gref shared_block ->
+    Shared_page_pool.use nf.t.tx_pool (fun ~id gref shared_block ->
       let len, datav = Cstruct.fillv ~src:datav ~dst:shared_block in
       (* [size] includes extra pages to follow later *)
       let size = match size with |None -> len |Some s -> s in
       Stats.tx nf.t.stats (Int64.of_int size);
-      let id = gref in
       let request = { TX.Request.
         id;
         gref = Int32.of_int gref;

lib/shared_page_pool.ml

@@ -18,24 +18,33 @@ open Lwt.Infix
 
 let return = Lwt.return
 
-type block = Gnt.gntref * Cstruct.t
+let max_pages = 256
+
+type block = {
+  id : Cstruct.uint16;
+  gref : Gnt.gntref;
+  data : Cstruct.t;
+}
 type t = {
   grant : Gnt.gntref -> Io_page.t -> unit;
+  mutable next_id : Cstruct.uint16;
   mutable blocks : block list;
   mutable in_use : int;
   mutable shutdown : bool;
+  avail : unit Lwt_condition.t;  (* Fires when free list becomes non-empty *)
 }
 
 let page_size = Io_page.round_to_page_size 1
 let block_size = page_size / 2
 
-let make grant = { grant; blocks = []; shutdown = false; in_use = 0 }
+let make grant = { next_id = 0; grant; blocks = []; shutdown = false; in_use = 0; avail = Lwt_condition.create () }
 
 let shutdown t =
   t.shutdown <- true;
+  Lwt_condition.broadcast t.avail ();   (* Wake anyone who's still waiting for free pages *)
   if t.in_use = 0 then (
-    t.blocks |> List.iter (fun (gref, block) ->
-      if block.Cstruct.off = 0 then (
+    t.blocks |> List.iter (fun {id = _; gref; data} ->
+      if data.Cstruct.off = 0 then (
         Gnt.Gntshr.end_access gref;
         Gnt.Gntshr.put gref;
       )
@@ -52,33 +61,44 @@ let alloc t =
   return (gnt, Io_page.to_cstruct page)
 
 let put t block =
+  let was_empty = (t.blocks = []) in
   t.blocks <- block :: t.blocks;
   t.in_use <- t.in_use - 1;
+  if was_empty then Lwt_condition.broadcast t.avail ();
   if t.in_use = 0 && t.shutdown then shutdown t
 
-let use t fn =
+let use_block t fn block =
+  let {id; gref; data} = block in
+  t.in_use <- t.in_use + 1;
+  Lwt.try_bind
+    (fun () -> fn ~id gref data)
+    (fun (_, release as result) ->
+      Lwt.on_termination release (fun () -> put t block);
+      return result
+    )
+    (fun ex -> put t block; Lwt.fail ex)
+
+let rec use t fn =
   if t.shutdown then
     failwith "Shared_page_pool.use after shutdown";
-  begin match t.blocks with
+  match t.blocks with
+  | [] when t.next_id >= max_pages ->
+      MProf.Trace.label "Shared_page_pool waiting for free";
+      Lwt_condition.wait t.avail >>= fun () -> use t fn
   | [] ->
       (* Frames normally fit within 2048 bytes, so we split each page in half. *)
-      alloc t >>= fun (gntref, page) ->
+      alloc t >>= fun (gref, page) ->
       let b1 = Cstruct.sub page 0 block_size in
       let b2 = Cstruct.shift page block_size in
-      t.blocks <- (gntref, b2) :: t.blocks;
-      return (gntref, b1)
+      let id1 = t.next_id in
+      let id2 = t.next_id + 1 in
+      t.next_id <- t.next_id + 2;
+      t.blocks <- {id = id2; gref; data = b2} :: t.blocks;
+      Lwt_condition.broadcast t.avail ();
+      use_block t fn {id = id1; gref; data = b1}
   | hd :: tl ->
       t.blocks <- tl;
-      return hd
-  end >>= fun (gntref, block as grant) ->
-  t.in_use <- t.in_use + 1;
-  Lwt.try_bind
-    (fun () -> fn gntref block)
-    (fun (_, release as result) ->
-      Lwt.on_termination release (fun () -> put t grant);
-      return result
-    )
-    (fun ex -> put t grant; Lwt.fail ex)
+      use_block t fn hd
 
 let blocks_needed bytes =
   (bytes + block_size - 1) / block_size

lib/shared_page_pool.mli

@@ -24,8 +24,9 @@ val make : (Gnt.gntref -> Io_page.t -> unit) -> t
 (** [make grant_access] is a shared pool of blocks. When a new page is first
     allocated, [grant_access] is called to share it. *)
 
-val use : t -> (Gnt.gntref -> Cstruct.t -> ('a * unit Lwt.t) Lwt.t) -> ('a * unit Lwt.t) Lwt.t
-(** [use t fn] calls [fn gref block] with a free shared block of memory.
+val use : t -> (id:Cstruct.uint16 -> Gnt.gntref -> Cstruct.t -> ('a * unit Lwt.t) Lwt.t) -> ('a * unit Lwt.t) Lwt.t
+(** [use t fn] calls [fn ~id gref block] with a free shared block of memory
+    and that block's unique ID (note: the [gref] is NOT unique).
     The function should return a thread that indicates when the request has
     been added to the queue, by returning a result value and a second thread
     indicating when the block can be returned to the pool. *)