Source file quota.ml

(*
 * Copyright (C) Citrix Systems Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation; version 2.1 only. with the special
 * exception on linking described in file LICENSE.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *)

let warn fmt = Logging.warn "quota" fmt

exception Limit_reached
exception Data_too_big
exception Transaction_opened

type domid = int

(* Global defaults *)
let maxent = ref 10000
let maxsize = ref 4096
let maxwatch = ref 50
let maxtransaction = ref 20
let maxwatchevent = ref 256

type overrides = (int, int) Hashtbl.t

let maxent_overrides = Hashtbl.create 10
let maxwatch_overrides = Hashtbl.create 10
let maxtransaction_overrides = Hashtbl.create 10
let maxwatchevent_overrides = Hashtbl.create 10

let get_override t domid =
  if Hashtbl.mem t domid then Some (Hashtbl.find t domid) else None

let set_override t domid override =
  match override with
  | None -> Hashtbl.remove t domid
  | Some override -> Hashtbl.replace t domid override

let list_overrides t = Hashtbl.fold (fun domid x acc -> (domid, x) :: acc) t []

let of_domain t default domid =
  if Hashtbl.mem t domid then Hashtbl.find t domid else !default

let maxent_of_domain = of_domain maxent_overrides maxent
let maxwatch_of_domain = of_domain maxwatch_overrides maxwatch
let maxtransaction_of_domain = of_domain maxtransaction_overrides maxtransaction
let maxwatchevent_of_domain = of_domain maxwatchevent_overrides maxwatchevent

type t = { cur : (domid, int) Hashtbl.t (* current domains entry usage *) }

let create () = { cur = Hashtbl.create 100 }
let copy quota = { cur = Hashtbl.copy quota.cur }

(*let del quota id = Hashtbl.remove quota.cur id*)

let check _quota id size =
  if size > !maxsize then (
    warn "domain %u err create entry: data too big %d" id size;
    raise Data_too_big)

let list quota =
  Hashtbl.fold (fun domid x acc -> (domid, x) :: acc) quota.cur []

let get quota id =
  if Hashtbl.mem quota.cur id then Hashtbl.find quota.cur id else 0

let set quota id nb =
  if nb = 0 then Hashtbl.remove quota.cur id
  else if Hashtbl.mem quota.cur id then Hashtbl.replace quota.cur id nb
  else Hashtbl.add quota.cur id nb

let decr quota id =
  let nb = get quota id in
  if nb > 0 then set quota id (nb - 1)

let incr quota id =
  let nb = get quota id in
  let maxent = maxent_of_domain id in
  if nb >= maxent then raise Limit_reached;
  set quota id (nb + 1)

let union quota diff =
  Hashtbl.iter (fun id nb -> set quota id (get quota id + nb)) diff.cur

let merge orig_quota mod_quota dest_quota =
  Hashtbl.iter
    (fun id nb ->
      let diff = nb - get orig_quota id in
      if diff <> 0 then set dest_quota id (get dest_quota id + diff))
    mod_quota.cur