<commit_before>(*
 * 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.
 *)

open Xenops_interface
open Xenops_server_plugin
open Xenops_utils
open Xenops_task

type context = {
	transferred_fd: Unix.file_descr option;
	(** some API calls take a file descriptor argument *)
}

let make_context () = {
	transferred_fd = None
}

let instance_id = Uuidm.to_string (Uuidm.create `V4)

let query _ _ _ = {
	Query.name = "xenops";
	vendor = "XCP";
	version = "0.1";
	features = [];
	instance_id = instance_id;
}

let backend = ref None
let get_backend () = match !backend with
  | Some x -> x 
  | None -> failwith "No backend implementation set"

let ignore_exception msg f x =
	try f x
	with
		| e ->
			debug "%s: ignoring exception: %s" msg (e |> exnty_of_exn |> Exception.rpc_of_exnty |> Jsonrpc.to_string)

let filter_prefix prefix xs =
	List.filter_map
		(fun x ->
			if String.startswith prefix x
			then Some (String.sub x (String.length prefix) (String.length x - (String.length prefix)))
			else None) xs

type atomic =
	| VBD_eject of Vbd.id
	| VIF_plug of Vif.id
	| VIF_unplug of Vif.id * bool
	| VIF_move of Vif.id * Network.t
	| VIF_set_carrier of Vif.id * bool
	| VIF_set_locking_mode of Vif.id * Vif.locking_mode
	| VM_hook_script of (Vm.id * Xenops_hooks.script * string)
	| VBD_plug of Vbd.id
	| VBD_epoch_begin of Vbd.id
	| VBD_epoch_end of Vbd.id
	| VBD_set_qos of Vbd.id
	| VBD_unplug of Vbd.id * bool
	| VBD_insert of Vbd.id * disk
	| VM_remove of Vm.id
	| PCI_plug of Pci.id
	| PCI_unplug of Pci.id
	| VM_set_xsdata of (Vm.id * (string * string) list)
	| VM_set_vcpus of (Vm.id * int)
	| VM_set_shadow_multiplier of (Vm.id * float)
	| VM_set_memory_dynamic_range of (Vm.id * int64 * int64)
	| VM_pause of Vm.id
	| VM_unpause of Vm.id
	| VM_set_domain_action_request of (Vm.id * domain_action_request option)
	| VM_create_device_model of (Vm.id * bool)
	| VM_destroy_device_model of Vm.id
	| VM_destroy of Vm.id
	| VM_create of (Vm.id * int64 option)
	| VM_build of Vm.id
	| VM_shutdown_domain of (Vm.id * shutdown_request * float)
	| VM_s3suspend of Vm.id
	| VM_s3resume of Vm.id
	| VM_save of (Vm.id * flag list * data)
	| VM_restore of (Vm.id * data)
	| VM_delay of (Vm.id * float) (** used to suppress fast reboot loops *)

with rpc

let string_of_atomic x = x |> rpc_of_atomic |> Jsonrpc.to_string

type operation =
	| VM_start of Vm.id
	| VM_poweroff of (Vm.id * float option)
	| VM_shutdown of (Vm.id * float option)
	| VM_reboot of (Vm.id * float option)
	| VM_suspend of (Vm.id * data)
	| VM_resume of (Vm.id * data)
	| VM_restore_devices of Vm.id
	| VM_migrate of (Vm.id * (string * string) list * (string * Network.t) list * string)
	| VM_receive_memory of (Vm.id * int64 * Unix.file_descr)
	| VM_check_state of Vm.id
	| PCI_check_state of Pci.id
	| VBD_check_state of Vbd.id
	| VIF_check_state of Vif.id
	| Atomic of atomic
with rpc

let string_of_operation x = x |> rpc_of_operation |> Jsonrpc.to_string

module TASK = struct
	open Xenops_task

	let task x = {
		Task.id = x.id;
		debug_info = x.debug_info;
		ctime = x.ctime;
		state = x.state;
		subtasks = x.subtasks;
	}
	let cancel _ dbg id =
		Xenops_task.cancel tasks id
	let stat' id =
		Mutex.execute tasks.m
			(fun () ->
				find_locked tasks id |> task
			)
	let signal id =
		Mutex.execute tasks.m
			(fun () ->
				if exists_locked tasks id then begin
					debug "TASK.signal %s = %s" id ((find_locked tasks id).state |> Task.rpc_of_state |> Jsonrpc.to_string);
					Updates.add (Dynamic.Task id) updates
				end else debug "TASK.signal %s (object deleted)" id
			)
	let stat _ dbg id = stat' id 
	let destroy' id = destroy tasks id; Updates.remove (Dynamic.Task id) updates
	let destroy _ dbg id = destroy' id
	let list _ dbg = list tasks |> List.map task
end

module VM_DB = struct
	include TypedTable(struct
		include Vm
		let namespace = "VM"
		type key = id
		let key id = [ id; "config" ]
	end)
	let ids () : Vm.id list =
		list []
	let list () =
		debug "VM.list";
		let vms = ids () |> List.map read |> List.filter_map (fun x -> x) in
		let module B = (val get_backend () : S) in
		let states = List.map B.VM.get_state vms in
		List.combine vms states
	let m = Mutex.create ()
	let signal id =
		debug "VM_DB.signal %s" id;
		Mutex.execute m
			(fun () ->
				if exists id
				then Updates.add (Dynamic.Vm id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Vm id) updates;
				remove id
			)
end

module PCI_DB = struct
	include TypedTable(struct
		include Pci
		let namespace = "PCI"
		type key = id
		let key k = [ fst k; "pci." ^ (snd k) ]	
	end)
	let vm_of = fst
	let string_of_id (a, b) = a ^ "." ^ b

	let ids vm : Pci.id list =
		list [ vm ] |> (filter_prefix "pci.") |> List.map (fun id -> (vm, id))
	let pcis vm = ids vm |> List.map read |> dropnone
	let list vm =
		debug "PCI.list";
		let xs = pcis vm in
		let module B = (val get_backend () : S) in
		let states = List.map (B.PCI.get_state vm) xs in
		List.combine xs states
	let m = Mutex.create ()
	let signal id =
		debug "PCI_DB.signal %s" (string_of_id id);
		Mutex.execute m
			(fun () ->
				if exists id
				then Updates.add (Dynamic.Pci id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Pci id) updates;
				remove id
			)
end

module VBD_DB = struct
	include TypedTable(struct
		include Vbd
		let namespace = "VM"
		type key = id
		let key k = [ fst k; "vbd." ^ (snd k) ]
	end)
	let vm_of = fst
	let string_of_id (a, b) = a ^ "." ^ b

	let ids vm : Vbd.id list =
		list [ vm ] |> (filter_prefix "vbd.") |> List.map (fun id -> (vm, id))
	let vbds vm = ids vm |> List.map read |> dropnone
	let list vm =
		debug "VBD.list";
		let vbds' = vbds vm in
		let module B = (val get_backend () : S) in
		let states = List.map (B.VBD.get_state vm) vbds' in
		List.combine vbds' states
	let m = Mutex.create ()
	let signal id =
		debug "VBD_DB.signal %s" (string_of_id id);
		Mutex.execute m
			(fun () ->
				if exists id
				then Updates.add (Dynamic.Vbd id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Vbd id) updates;
				remove id
			)
end

module VIF_DB = struct
	include TypedTable(struct
		include Vif
		let namespace = "VM"
		type key = id
		let key k = [ fst k; "vif." ^ (snd k) ]
	end)
	let vm_of = fst
	let string_of_id (a, b) = a ^ "." ^ b

	let ids vm : Vif.id list =
		list [ vm ] |> (filter_prefix "vif.") |> List.map (fun id -> (vm, id))
	let vifs vm = ids vm |> List.map read |> dropnone
 	let list vm =
		let vifs' = vifs vm in
		let module B = (val get_backend () : S) in
		let states = List.map (B.VIF.get_state vm) vifs' in
		List.combine vifs' states
	let m = Mutex.create ()
	let signal id =
		debug "VIF_DB.signal %s" (string_of_id id);
		Mutex.execute m
			(fun () ->
				Updates.add (Dynamic.Vif id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Vif id) updates;
				remove id
			)
end

module StringMap = Map.Make(struct type t = string let compare = compare end)

let push_with_coalesce should_keep item queue =
	(* [filter_with_memory p xs] returns elements [x \in xs] where [p (x_i, [x_0...x_i-1])] *)
	let filter_with_memory p xs =
		List.fold_left (fun (acc, xs) x -> xs :: acc, x :: xs) ([], []) xs
		|> fst |> List.rev |> List.combine xs (* association list of (element, all previous elements) *)
		|> List.filter p
		|> List.map fst in

	let to_list queue = Queue.fold (fun xs x -> x :: xs) [] queue |> List.rev in
	let of_list xs =
		let q = Queue.create () in
		List.iter (fun x -> Queue.push x q) xs;
		q in

	Queue.push item queue;
	let queue' =
		to_list queue
	|> filter_with_memory (fun (this, prev) -> should_keep this prev)
	|> of_list in
	Queue.clear queue;
	Queue.transfer queue' queue

module Queues = struct
	(** A set of queues where 'pop' operates on each queue in a round-robin fashion *)

	type tag = string
	(** Each distinct 'tag' value creates a separate virtual queue *)

	type 'a t = {
		mutable qs: 'a Queue.t StringMap.t;
		mutable last_tag: string;
		m: Mutex.t;
		c: Condition.t;
	}

	let create () = {
 		qs = StringMap.empty;
		last_tag = "";
		m = Mutex.create ();
		c = Condition.create ();
	}

	let get tag qs =
		Mutex.execute qs.m
			(fun () ->
				if StringMap.mem tag qs.qs then StringMap.find tag qs.qs else Queue.create ()
			)

	let tags qs =
		Mutex.execute qs.m
			(fun () ->
				StringMap.fold (fun x _ acc -> x :: acc) qs.qs []
			)

	let get_last_tag qs =
		Mutex.execute qs.m
			(fun () ->
				qs.last_tag
			)

	let push_with_coalesce should_keep tag item qs =
		Mutex.execute qs.m
			(fun () ->
				let q = if StringMap.mem tag qs.qs then StringMap.find tag qs.qs else Queue.create () in
				push_with_coalesce should_keep item q;
				qs.qs <- StringMap.add tag q qs.qs;
				Condition.signal qs.c
			)

	let pop qs =
		Mutex.execute qs.m
			(fun () ->
				while StringMap.is_empty qs.qs do
					Condition.wait qs.c qs.m;
				done;
				(* partition based on last_tag *)
				let before, after = StringMap.partition (fun x _ -> x <= qs.last_tag) qs.qs in
				(* the min_binding in the 'after' is the next queue *)
				let last_tag, q = StringMap.min_binding (if StringMap.is_empty after then before else after) in
				qs.last_tag <- last_tag;
				let item = Queue.pop q in
				(* remove empty queues from the whole mapping *)
				qs.qs <- if Queue.is_empty q then StringMap.remove last_tag qs.qs else qs.qs;
				last_tag, item
			)

	let transfer_tag tag a b =
		Mutex.execute a.m
			(fun () ->
				Mutex.execute b.m
					(fun () ->
						if StringMap.mem tag a.qs then begin
							b.qs <- StringMap.add tag (StringMap.find tag a.qs) b.qs;
							a.qs <- StringMap.remove tag a.qs;
							Condition.signal b.c
						end
					)
			)
end

module Redirector = struct
	(* When a thread is not actively processing a queue, items are placed here: *)
	let default = Queues.create ()

	(* When a thread is actively processing a queue, items are redirected to a thread-private queue *)
	let overrides = ref StringMap.empty
	let m = Mutex.create ()

	let should_keep (op, _) prev = match op with
		| VM_check_state (_)
		| PCI_check_state (_)
		| VBD_check_state (_)
		| VIF_check_state (_) ->
			let prev' = List.map fst prev in
			not(List.mem op prev')
		| _ -> true

	let push tag item =
		Debug.with_thread_associated "queue"
			(fun () ->
				Mutex.execute m
					(fun () ->
						let q, redirected = if StringMap.mem tag !overrides then StringMap.find tag !overrides, true else default, false in
						debug "Queue.push %s onto %s%s:[ %s ]" (string_of_operation (fst item)) (if redirected then "redirected " else "") tag (String.concat ", " (List.rev (Queue.fold (fun acc (b, _) -> string_of_operation b :: acc) [] (Queues.get tag q))));

						Queues.push_with_coalesce should_keep tag item q
					)
			) ()

	let pop =
		(* We must prevent worker threads all calling Queues.pop before we've
		   successfully put the redirection in place. Otherwise we end up with
		   parallel threads operating on the same VM. *)
		let n = Mutex.create () in
		fun () ->
			Mutex.execute n
				(fun () ->
					let tag, item = Queues.pop default in
					Mutex.execute m
						(fun () ->
							let q = Queues.create () in
							Queues.transfer_tag tag default q;
							overrides := StringMap.add tag q !overrides;
							(* All items with [tag] will enter queue [q] *)
							tag, q, item
						)
				)

	let finished tag queue =
		Mutex.execute m
			(fun () ->
				Queues.transfer_tag tag queue default;
				overrides := StringMap.remove tag !overrides
				(* All items with [tag] will enter the default queue *)
			)

	module Dump = struct
		type q = {
			tag: string;
			items: operation list
		} with rpc
		type t = q list with rpc

		let make () =
			Mutex.execute m
				(fun () ->
					let one queue =
						List.map
							(fun t ->
							{ tag = t; items = List.rev (Queue.fold (fun acc (b, _) -> b :: acc) [] (Queues.get t queue)) }
							) (Queues.tags queue) in
					List.concat (List.map one (default :: (List.map snd (StringMap.bindings !overrides))))
				)

	end
end		

module Worker = struct
	type state =
		| Idle
		| Processing of (operation * Xenops_task.t)
		| Shutdown_requested
		| Shutdown
	type t = {
		mutable state: state;
		mutable shutdown_requested: bool;
		m: Mutex.t;
		c: Condition.t;
		mutable t: Thread.t option;
	}

	let get_state_locked t =
		if t.shutdown_requested
		then Shutdown_requested
		else t.state

	let get_state t =
		Mutex.execute t.m
			(fun () ->
				get_state_locked t
			)

	let join t =
		Mutex.execute t.m
			(fun () ->
				assert (t.state = Shutdown);
				Opt.iter Thread.join t.t
			)

	let is_active t =
		Mutex.execute t.m
			(fun () ->
				match get_state_locked t with
					| Idle | Processing (_, _) -> true
					| Shutdown_requested | Shutdown -> false
			)

	let shutdown t =
		Mutex.execute t.m
			(fun () ->
				if not t.shutdown_requested then begin
					t.shutdown_requested <- true;
					true (* success *)
				end else false
			)

	let restart t =
		Mutex.execute t.m
			(fun () ->
				if t.shutdown_requested && t.state <> Shutdown then begin
					t.shutdown_requested <- false;
					true (* success *)
				end else false
			)

	let create () =
		let t = {
			state = Idle;
			shutdown_requested = false;
			m = Mutex.create ();
			c = Condition.create ();
			t = None;
		} in
		let thread = Thread.create
			(fun () ->
				while not(Mutex.execute t.m (fun () ->
					if t.shutdown_requested then t.state <- Shutdown;
					t.shutdown_requested
				)) do
					Mutex.execute t.m (fun () -> t.state <- Idle);
					let tag, queue, (op, item) = Redirector.pop () in (* blocks here *)
					debug "Queue.pop returned %s" (string_of_operation op);
					Mutex.execute t.m (fun () -> t.state <- Processing (op, item));
					begin
						try
							Debug.with_thread_associated
								item.Xenops_task.debug_info
								(fun () ->
									debug "Task %s reference %s: %s" item.Xenops_task.id item.Xenops_task.debug_info (string_of_operation op);
									Xenops_task.run item
								) ()
						with e ->
							debug "Queue caught: %s" (Printexc.to_string e)
					end;
					Redirector.finished tag queue;
					(* The task must have succeeded or failed. *)
					begin match item.Xenops_task.state with
						| Task.Pending _ ->
							error "Task %s has been left in a Pending state" item.Xenops_task.id;
							item.Xenops_task.state <- Task.Failed (Internal_error "Task left in Pending state" |> exnty_of_exn |> Exception.rpc_of_exnty)
						| _ -> ()
					end;
					TASK.signal item.Xenops_task.id
				done
			) () in
		t.t <- Some thread;
		t
end

module WorkerPool = struct

	(* Store references to Worker.ts here *)
	let pool = ref []
	let m = Mutex.create ()


	module Dump = struct
		type task = {
			id: string;
			ctime: string;
			debug_info: string;
			subtasks: (string * string) list;
		} with rpc
		let of_task t = {
				id = t.Xenops_task.id;
				ctime = t.Xenops_task.ctime |> Date.of_float |> Date.to_string;
				debug_info = t.Xenops_task.debug_info;
				subtasks = List.map (fun (name, state) -> name, state |> Task.rpc_of_state |> Jsonrpc.to_string) t.Xenops_task.subtasks |> List.rev;
			}
		type w = {
			state: string;
			task: task option;
		} with rpc
		type t = w list with rpc
		let make () =
			Mutex.execute m
				(fun () ->
					List.map
						(fun t ->
							match Worker.get_state t with
								| Worker.Idle -> { state = "Idle"; task = None }
								| Worker.Processing (op, task) -> { state = Printf.sprintf "Processing %s" (string_of_operation op); task = Some (of_task task) }
								| Worker.Shutdown_requested -> { state = "Shutdown_requested"; task = None }
								| Worker.Shutdown -> { state = "Shutdown"; task = None }
						) !pool
				)
	end

	(* Compute the number of active threads ie those which will continue to operate *)
	let count_active () =
		Mutex.execute m
			(fun () ->
				List.map Worker.is_active !pool |> List.filter id |> List.length
			)

	let find_one f = List.fold_left (fun acc x -> acc || (f x)) false

	(* Clean up any shutdown threads and remove them from the master list *)
	let gc pool =
		List.fold_left
			(fun acc w ->
				if Worker.get_state w = Worker.Shutdown then begin
					Worker.join w;
					acc
				end else w :: acc) [] pool

	let incr () =
		debug "Adding a new worker to the thread pool";
		Mutex.execute m
			(fun () ->
				pool := gc !pool;
				if not(find_one Worker.restart !pool)
				then pool := (Worker.create ()) :: !pool
			)

	let decr () =
		debug "Removing a worker from the thread pool";
		Mutex.execute m
			(fun () -><commit_msg>Inline 'id' (let id x = x)<commit_after>(*
 * 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.
 *)

open Xenops_interface
open Xenops_server_plugin
open Xenops_utils
open Xenops_task

type context = {
	transferred_fd: Unix.file_descr option;
	(** some API calls take a file descriptor argument *)
}

let make_context () = {
	transferred_fd = None
}

let instance_id = Uuidm.to_string (Uuidm.create `V4)

let query _ _ _ = {
	Query.name = "xenops";
	vendor = "XCP";
	version = "0.1";
	features = [];
	instance_id = instance_id;
}

let backend = ref None
let get_backend () = match !backend with
  | Some x -> x 
  | None -> failwith "No backend implementation set"

let ignore_exception msg f x =
	try f x
	with
		| e ->
			debug "%s: ignoring exception: %s" msg (e |> exnty_of_exn |> Exception.rpc_of_exnty |> Jsonrpc.to_string)

let filter_prefix prefix xs =
	List.filter_map
		(fun x ->
			if String.startswith prefix x
			then Some (String.sub x (String.length prefix) (String.length x - (String.length prefix)))
			else None) xs

type atomic =
	| VBD_eject of Vbd.id
	| VIF_plug of Vif.id
	| VIF_unplug of Vif.id * bool
	| VIF_move of Vif.id * Network.t
	| VIF_set_carrier of Vif.id * bool
	| VIF_set_locking_mode of Vif.id * Vif.locking_mode
	| VM_hook_script of (Vm.id * Xenops_hooks.script * string)
	| VBD_plug of Vbd.id
	| VBD_epoch_begin of Vbd.id
	| VBD_epoch_end of Vbd.id
	| VBD_set_qos of Vbd.id
	| VBD_unplug of Vbd.id * bool
	| VBD_insert of Vbd.id * disk
	| VM_remove of Vm.id
	| PCI_plug of Pci.id
	| PCI_unplug of Pci.id
	| VM_set_xsdata of (Vm.id * (string * string) list)
	| VM_set_vcpus of (Vm.id * int)
	| VM_set_shadow_multiplier of (Vm.id * float)
	| VM_set_memory_dynamic_range of (Vm.id * int64 * int64)
	| VM_pause of Vm.id
	| VM_unpause of Vm.id
	| VM_set_domain_action_request of (Vm.id * domain_action_request option)
	| VM_create_device_model of (Vm.id * bool)
	| VM_destroy_device_model of Vm.id
	| VM_destroy of Vm.id
	| VM_create of (Vm.id * int64 option)
	| VM_build of Vm.id
	| VM_shutdown_domain of (Vm.id * shutdown_request * float)
	| VM_s3suspend of Vm.id
	| VM_s3resume of Vm.id
	| VM_save of (Vm.id * flag list * data)
	| VM_restore of (Vm.id * data)
	| VM_delay of (Vm.id * float) (** used to suppress fast reboot loops *)

with rpc

let string_of_atomic x = x |> rpc_of_atomic |> Jsonrpc.to_string

type operation =
	| VM_start of Vm.id
	| VM_poweroff of (Vm.id * float option)
	| VM_shutdown of (Vm.id * float option)
	| VM_reboot of (Vm.id * float option)
	| VM_suspend of (Vm.id * data)
	| VM_resume of (Vm.id * data)
	| VM_restore_devices of Vm.id
	| VM_migrate of (Vm.id * (string * string) list * (string * Network.t) list * string)
	| VM_receive_memory of (Vm.id * int64 * Unix.file_descr)
	| VM_check_state of Vm.id
	| PCI_check_state of Pci.id
	| VBD_check_state of Vbd.id
	| VIF_check_state of Vif.id
	| Atomic of atomic
with rpc

let string_of_operation x = x |> rpc_of_operation |> Jsonrpc.to_string

module TASK = struct
	open Xenops_task

	let task x = {
		Task.id = x.id;
		debug_info = x.debug_info;
		ctime = x.ctime;
		state = x.state;
		subtasks = x.subtasks;
	}
	let cancel _ dbg id =
		Xenops_task.cancel tasks id
	let stat' id =
		Mutex.execute tasks.m
			(fun () ->
				find_locked tasks id |> task
			)
	let signal id =
		Mutex.execute tasks.m
			(fun () ->
				if exists_locked tasks id then begin
					debug "TASK.signal %s = %s" id ((find_locked tasks id).state |> Task.rpc_of_state |> Jsonrpc.to_string);
					Updates.add (Dynamic.Task id) updates
				end else debug "TASK.signal %s (object deleted)" id
			)
	let stat _ dbg id = stat' id 
	let destroy' id = destroy tasks id; Updates.remove (Dynamic.Task id) updates
	let destroy _ dbg id = destroy' id
	let list _ dbg = list tasks |> List.map task
end

module VM_DB = struct
	include TypedTable(struct
		include Vm
		let namespace = "VM"
		type key = id
		let key id = [ id; "config" ]
	end)
	let ids () : Vm.id list =
		list []
	let list () =
		debug "VM.list";
		let vms = ids () |> List.map read |> List.filter_map (fun x -> x) in
		let module B = (val get_backend () : S) in
		let states = List.map B.VM.get_state vms in
		List.combine vms states
	let m = Mutex.create ()
	let signal id =
		debug "VM_DB.signal %s" id;
		Mutex.execute m
			(fun () ->
				if exists id
				then Updates.add (Dynamic.Vm id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Vm id) updates;
				remove id
			)
end

module PCI_DB = struct
	include TypedTable(struct
		include Pci
		let namespace = "PCI"
		type key = id
		let key k = [ fst k; "pci." ^ (snd k) ]	
	end)
	let vm_of = fst
	let string_of_id (a, b) = a ^ "." ^ b

	let ids vm : Pci.id list =
		list [ vm ] |> (filter_prefix "pci.") |> List.map (fun id -> (vm, id))
	let pcis vm = ids vm |> List.map read |> dropnone
	let list vm =
		debug "PCI.list";
		let xs = pcis vm in
		let module B = (val get_backend () : S) in
		let states = List.map (B.PCI.get_state vm) xs in
		List.combine xs states
	let m = Mutex.create ()
	let signal id =
		debug "PCI_DB.signal %s" (string_of_id id);
		Mutex.execute m
			(fun () ->
				if exists id
				then Updates.add (Dynamic.Pci id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Pci id) updates;
				remove id
			)
end

module VBD_DB = struct
	include TypedTable(struct
		include Vbd
		let namespace = "VM"
		type key = id
		let key k = [ fst k; "vbd." ^ (snd k) ]
	end)
	let vm_of = fst
	let string_of_id (a, b) = a ^ "." ^ b

	let ids vm : Vbd.id list =
		list [ vm ] |> (filter_prefix "vbd.") |> List.map (fun id -> (vm, id))
	let vbds vm = ids vm |> List.map read |> dropnone
	let list vm =
		debug "VBD.list";
		let vbds' = vbds vm in
		let module B = (val get_backend () : S) in
		let states = List.map (B.VBD.get_state vm) vbds' in
		List.combine vbds' states
	let m = Mutex.create ()
	let signal id =
		debug "VBD_DB.signal %s" (string_of_id id);
		Mutex.execute m
			(fun () ->
				if exists id
				then Updates.add (Dynamic.Vbd id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Vbd id) updates;
				remove id
			)
end

module VIF_DB = struct
	include TypedTable(struct
		include Vif
		let namespace = "VM"
		type key = id
		let key k = [ fst k; "vif." ^ (snd k) ]
	end)
	let vm_of = fst
	let string_of_id (a, b) = a ^ "." ^ b

	let ids vm : Vif.id list =
		list [ vm ] |> (filter_prefix "vif.") |> List.map (fun id -> (vm, id))
	let vifs vm = ids vm |> List.map read |> dropnone
 	let list vm =
		let vifs' = vifs vm in
		let module B = (val get_backend () : S) in
		let states = List.map (B.VIF.get_state vm) vifs' in
		List.combine vifs' states
	let m = Mutex.create ()
	let signal id =
		debug "VIF_DB.signal %s" (string_of_id id);
		Mutex.execute m
			(fun () ->
				Updates.add (Dynamic.Vif id) updates
			)
	let remove id =
		Mutex.execute m
			(fun () ->
				Updates.remove (Dynamic.Vif id) updates;
				remove id
			)
end

module StringMap = Map.Make(struct type t = string let compare = compare end)

let push_with_coalesce should_keep item queue =
	(* [filter_with_memory p xs] returns elements [x \in xs] where [p (x_i, [x_0...x_i-1])] *)
	let filter_with_memory p xs =
		List.fold_left (fun (acc, xs) x -> xs :: acc, x :: xs) ([], []) xs
		|> fst |> List.rev |> List.combine xs (* association list of (element, all previous elements) *)
		|> List.filter p
		|> List.map fst in

	let to_list queue = Queue.fold (fun xs x -> x :: xs) [] queue |> List.rev in
	let of_list xs =
		let q = Queue.create () in
		List.iter (fun x -> Queue.push x q) xs;
		q in

	Queue.push item queue;
	let queue' =
		to_list queue
	|> filter_with_memory (fun (this, prev) -> should_keep this prev)
	|> of_list in
	Queue.clear queue;
	Queue.transfer queue' queue

module Queues = struct
	(** A set of queues where 'pop' operates on each queue in a round-robin fashion *)

	type tag = string
	(** Each distinct 'tag' value creates a separate virtual queue *)

	type 'a t = {
		mutable qs: 'a Queue.t StringMap.t;
		mutable last_tag: string;
		m: Mutex.t;
		c: Condition.t;
	}

	let create () = {
 		qs = StringMap.empty;
		last_tag = "";
		m = Mutex.create ();
		c = Condition.create ();
	}

	let get tag qs =
		Mutex.execute qs.m
			(fun () ->
				if StringMap.mem tag qs.qs then StringMap.find tag qs.qs else Queue.create ()
			)

	let tags qs =
		Mutex.execute qs.m
			(fun () ->
				StringMap.fold (fun x _ acc -> x :: acc) qs.qs []
			)

	let get_last_tag qs =
		Mutex.execute qs.m
			(fun () ->
				qs.last_tag
			)

	let push_with_coalesce should_keep tag item qs =
		Mutex.execute qs.m
			(fun () ->
				let q = if StringMap.mem tag qs.qs then StringMap.find tag qs.qs else Queue.create () in
				push_with_coalesce should_keep item q;
				qs.qs <- StringMap.add tag q qs.qs;
				Condition.signal qs.c
			)

	let pop qs =
		Mutex.execute qs.m
			(fun () ->
				while StringMap.is_empty qs.qs do
					Condition.wait qs.c qs.m;
				done;
				(* partition based on last_tag *)
				let before, after = StringMap.partition (fun x _ -> x <= qs.last_tag) qs.qs in
				(* the min_binding in the 'after' is the next queue *)
				let last_tag, q = StringMap.min_binding (if StringMap.is_empty after then before else after) in
				qs.last_tag <- last_tag;
				let item = Queue.pop q in
				(* remove empty queues from the whole mapping *)
				qs.qs <- if Queue.is_empty q then StringMap.remove last_tag qs.qs else qs.qs;
				last_tag, item
			)

	let transfer_tag tag a b =
		Mutex.execute a.m
			(fun () ->
				Mutex.execute b.m
					(fun () ->
						if StringMap.mem tag a.qs then begin
							b.qs <- StringMap.add tag (StringMap.find tag a.qs) b.qs;
							a.qs <- StringMap.remove tag a.qs;
							Condition.signal b.c
						end
					)
			)
end

module Redirector = struct
	(* When a thread is not actively processing a queue, items are placed here: *)
	let default = Queues.create ()

	(* When a thread is actively processing a queue, items are redirected to a thread-private queue *)
	let overrides = ref StringMap.empty
	let m = Mutex.create ()

	let should_keep (op, _) prev = match op with
		| VM_check_state (_)
		| PCI_check_state (_)
		| VBD_check_state (_)
		| VIF_check_state (_) ->
			let prev' = List.map fst prev in
			not(List.mem op prev')
		| _ -> true

	let push tag item =
		Debug.with_thread_associated "queue"
			(fun () ->
				Mutex.execute m
					(fun () ->
						let q, redirected = if StringMap.mem tag !overrides then StringMap.find tag !overrides, true else default, false in
						debug "Queue.push %s onto %s%s:[ %s ]" (string_of_operation (fst item)) (if redirected then "redirected " else "") tag (String.concat ", " (List.rev (Queue.fold (fun acc (b, _) -> string_of_operation b :: acc) [] (Queues.get tag q))));

						Queues.push_with_coalesce should_keep tag item q
					)
			) ()

	let pop =
		(* We must prevent worker threads all calling Queues.pop before we've
		   successfully put the redirection in place. Otherwise we end up with
		   parallel threads operating on the same VM. *)
		let n = Mutex.create () in
		fun () ->
			Mutex.execute n
				(fun () ->
					let tag, item = Queues.pop default in
					Mutex.execute m
						(fun () ->
							let q = Queues.create () in
							Queues.transfer_tag tag default q;
							overrides := StringMap.add tag q !overrides;
							(* All items with [tag] will enter queue [q] *)
							tag, q, item
						)
				)

	let finished tag queue =
		Mutex.execute m
			(fun () ->
				Queues.transfer_tag tag queue default;
				overrides := StringMap.remove tag !overrides
				(* All items with [tag] will enter the default queue *)
			)

	module Dump = struct
		type q = {
			tag: string;
			items: operation list
		} with rpc
		type t = q list with rpc

		let make () =
			Mutex.execute m
				(fun () ->
					let one queue =
						List.map
							(fun t ->
							{ tag = t; items = List.rev (Queue.fold (fun acc (b, _) -> b :: acc) [] (Queues.get t queue)) }
							) (Queues.tags queue) in
					List.concat (List.map one (default :: (List.map snd (StringMap.bindings !overrides))))
				)

	end
end		

module Worker = struct
	type state =
		| Idle
		| Processing of (operation * Xenops_task.t)
		| Shutdown_requested
		| Shutdown
	type t = {
		mutable state: state;
		mutable shutdown_requested: bool;
		m: Mutex.t;
		c: Condition.t;
		mutable t: Thread.t option;
	}

	let get_state_locked t =
		if t.shutdown_requested
		then Shutdown_requested
		else t.state

	let get_state t =
		Mutex.execute t.m
			(fun () ->
				get_state_locked t
			)

	let join t =
		Mutex.execute t.m
			(fun () ->
				assert (t.state = Shutdown);
				Opt.iter Thread.join t.t
			)

	let is_active t =
		Mutex.execute t.m
			(fun () ->
				match get_state_locked t with
					| Idle | Processing (_, _) -> true
					| Shutdown_requested | Shutdown -> false
			)

	let shutdown t =
		Mutex.execute t.m
			(fun () ->
				if not t.shutdown_requested then begin
					t.shutdown_requested <- true;
					true (* success *)
				end else false
			)

	let restart t =
		Mutex.execute t.m
			(fun () ->
				if t.shutdown_requested && t.state <> Shutdown then begin
					t.shutdown_requested <- false;
					true (* success *)
				end else false
			)

	let create () =
		let t = {
			state = Idle;
			shutdown_requested = false;
			m = Mutex.create ();
			c = Condition.create ();
			t = None;
		} in
		let thread = Thread.create
			(fun () ->
				while not(Mutex.execute t.m (fun () ->
					if t.shutdown_requested then t.state <- Shutdown;
					t.shutdown_requested
				)) do
					Mutex.execute t.m (fun () -> t.state <- Idle);
					let tag, queue, (op, item) = Redirector.pop () in (* blocks here *)
					debug "Queue.pop returned %s" (string_of_operation op);
					Mutex.execute t.m (fun () -> t.state <- Processing (op, item));
					begin
						try
							Debug.with_thread_associated
								item.Xenops_task.debug_info
								(fun () ->
									debug "Task %s reference %s: %s" item.Xenops_task.id item.Xenops_task.debug_info (string_of_operation op);
									Xenops_task.run item
								) ()
						with e ->
							debug "Queue caught: %s" (Printexc.to_string e)
					end;
					Redirector.finished tag queue;
					(* The task must have succeeded or failed. *)
					begin match item.Xenops_task.state with
						| Task.Pending _ ->
							error "Task %s has been left in a Pending state" item.Xenops_task.id;
							item.Xenops_task.state <- Task.Failed (Internal_error "Task left in Pending state" |> exnty_of_exn |> Exception.rpc_of_exnty)
						| _ -> ()
					end;
					TASK.signal item.Xenops_task.id
				done
			) () in
		t.t <- Some thread;
		t
end

module WorkerPool = struct

	(* Store references to Worker.ts here *)
	let pool = ref []
	let m = Mutex.create ()


	module Dump = struct
		type task = {
			id: string;
			ctime: string;
			debug_info: string;
			subtasks: (string * string) list;
		} with rpc
		let of_task t = {
				id = t.Xenops_task.id;
				ctime = t.Xenops_task.ctime |> Date.of_float |> Date.to_string;
				debug_info = t.Xenops_task.debug_info;
				subtasks = List.map (fun (name, state) -> name, state |> Task.rpc_of_state |> Jsonrpc.to_string) t.Xenops_task.subtasks |> List.rev;
			}
		type w = {
			state: string;
			task: task option;
		} with rpc
		type t = w list with rpc
		let make () =
			Mutex.execute m
				(fun () ->
					List.map
						(fun t ->
							match Worker.get_state t with
								| Worker.Idle -> { state = "Idle"; task = None }
								| Worker.Processing (op, task) -> { state = Printf.sprintf "Processing %s" (string_of_operation op); task = Some (of_task task) }
								| Worker.Shutdown_requested -> { state = "Shutdown_requested"; task = None }
								| Worker.Shutdown -> { state = "Shutdown"; task = None }
						) !pool
				)
	end

	(* Compute the number of active threads ie those which will continue to operate *)
	let count_active () =
		Mutex.execute m
			(fun () ->
				List.map Worker.is_active !pool |> List.filter (fun x -> x) |> List.length
			)

	let find_one f = List.fold_left (fun acc x -> acc || (f x)) false

	(* Clean up any shutdown threads and remove them from the master list *)
	let gc pool =
		List.fold_left
			(fun acc w ->
				if Worker.get_state w = Worker.Shutdown then begin
					Worker.join w;
					acc
				end else w :: acc) [] pool

	let incr () =
		debug "Adding a new worker to the thread pool";
		Mutex.execute m
			(fun () ->
				pool := gc !pool;
				if not(find_one Worker.restart !pool)
				then pool := (Worker.create ()) :: !pool
			)

	let decr () =
		debug "Removing a worker from the thread pool";
		Mutex.execute m
			(fun () ->
{'fraction_non_alphanumeric': 0.11542616871550793, 'fraction_numerical': 0.0009927876894326511, 'mean_word_length': 2.6909149692854832, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 230, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4782934', 'n_tokens_mistral': 14336, 'n_tokens_neox': 12497, 'n_words': 3989}