Build the segment tree! Commit this quick before something bad happens ...
[virt-df.git] / lib / diskimage.ml
1 (* Diskimage library for reading disk images.
2    (C) Copyright 2007-2008 Richard W.M. Jones, Red Hat Inc.
3    http://libvirt.org/
4
5    This program is free software; you can redistribute it and/or modify
6    it under the terms of the GNU General Public License as published by
7    the Free Software Foundation; either version 2 of the License, or
8    (at your option) any later version.
9
10    This program is distributed in the hope that it will be useful,
11    but WITHOUT ANY WARRANTY; without even the implied warranty of
12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13    GNU General Public License for more details.
14
15    You should have received a copy of the GNU General Public License
16    along with this program; if not, write to the Free Software
17    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
18  *)
19
20 open Printf
21 open ExtList
22 open Unix
23
24 open Int63.Operators
25
26 include Diskimage_utils
27
28 (* Use as the natural block size for disk images, but really we should
29  * use the 'blockdev -getbsz' command to find the real block size.
30  *)
31 let disk_block_size = ~^512
32
33 (*----------------------------------------------------------------------*)
34 (* The plug-ins. *)
35 let partition_types = [
36   Diskimage_mbr.plugin_id,
37     ("MBR", Diskimage_mbr.probe);
38 ]
39
40 let filesystem_types = [
41   Diskimage_ext2.plugin_id,
42     ("Linux ext2/3", Diskimage_ext2.probe);
43   Diskimage_linux_swap.plugin_id,
44     ("Linux swap", Diskimage_linux_swap.probe);
45   Diskimage_linux_swsuspend.plugin_id,
46     ("Linux s/w suspend", Diskimage_linux_swsuspend.probe);
47 ]
48
49 let lvm_types = [
50   Diskimage_lvm2.plugin_id,
51     ("Linux LVM2", Diskimage_lvm2.probe, Diskimage_lvm2.list);
52 ]
53
54 let name_of_parts id =
55   let name, _ = List.assoc id partition_types in
56   name
57 let name_of_filesystem id =
58   let name, _ = List.assoc id filesystem_types in
59   name
60 let name_of_lvm id =
61   let name, _, _ = List.assoc id lvm_types in
62   name
63
64 (* Probe a device for partitions.  Returns [Some parts] or [None]. *)
65 let probe_for_partitions dev =
66   if !debug then eprintf "probing for partitions on %s ...\n%!" dev#name;
67   let rec loop = function
68     | [] -> None
69     | (parts_plugin_id, (_, probe_fn)) :: rest ->
70         try Some (probe_fn dev)
71         with Not_found -> loop rest
72   in
73   let r = loop partition_types in
74   if !debug then (
75     match r with
76     | None -> eprintf "no partitions found on %s\n%!" dev#name
77     | Some { parts_plugin_id = name; parts = parts } ->
78         eprintf "found %d %s partitions on %s\n"
79           (List.length parts) name dev#name
80   );
81   r
82
83 (* Probe a device for a filesystem.  Returns [Some fs] or [None]. *)
84 let probe_for_filesystem dev =
85   if !debug then eprintf "probing for a filesystem on %s ...\n%!" dev#name;
86   let rec loop = function
87     | [] -> None
88     | (fs_name, (_, probe_fn)) :: rest ->
89         try Some (probe_fn dev)
90         with Not_found -> loop rest
91   in
92   let r = loop filesystem_types in
93   if !debug then (
94     match r with
95     | None -> eprintf "no filesystem found on %s\n%!" dev#name
96     | Some fs ->
97         eprintf "found a filesystem on %s:\n" dev#name;
98         eprintf "\t%s\n%!" fs.fs_plugin_id
99   );
100   r
101
102 (* Probe a device for a PV.  Returns [Some lvm_name] or [None]. *)
103 let probe_for_pv dev =
104   if !debug then eprintf "probing if %s is a PV ...\n%!" dev#name;
105   let rec loop = function
106     | [] -> None
107     | (lvm_name, (_, probe_fn, _)) :: rest ->
108         try Some (probe_fn lvm_name dev)
109         with Not_found -> loop rest
110   in
111   let r = loop lvm_types in
112   if !debug then (
113     match r with
114     | None -> eprintf "no PV found on %s\n%!" dev#name
115     | Some { lvm_plugin_id = name } ->
116         eprintf "%s contains a %s PV\n%!" dev#name name
117   );
118   r
119
120 let list_lvs lvm_name devs =
121   let _, _, list_lvs_fn = List.assoc lvm_name lvm_types in
122   list_lvs_fn devs
123
124 (*----------------------------------------------------------------------*)
125 (* Create machine description. *)
126 let open_machine name disks =
127   let disks = List.map (
128     fun (name, path) ->
129       let dev = new block_device path disk_block_size (* XXX *) in
130       { d_name = name; d_dev = dev; d_content = `Unknown }
131   ) disks in
132   { m_name = name; m_disks = disks; m_lv_filesystems = [] }
133
134 let close_machine { m_disks = m_disks } =
135   (* Only close the disks, assume all other devices are derived from them. *)
136   List.iter (fun { d_dev = d_dev } -> d_dev#close ()) m_disks
137
138 (* Main scanning function for filesystems. *)
139 let scan_machine ({ m_disks = m_disks } as machine) =
140   let m_disks = List.map (
141     fun ({ d_dev = dev } as disk) ->
142       let dev = (dev :> device) in
143       (* See if it is partitioned first. *)
144       let parts = probe_for_partitions dev in
145       match parts with
146       | Some parts ->
147           { disk with d_content = `Partitions parts }
148       | None ->
149           (* Not partitioned.  Does it contain a filesystem? *)
150           let fs = probe_for_filesystem dev in
151           match fs with
152           | Some fs ->
153               { disk with d_content = `Filesystem fs }
154           | None ->
155               (* Not partitioned, no filesystem, is it a PV? *)
156               let pv = probe_for_pv dev in
157               match pv with
158               | Some lvm_name ->
159                   { disk with d_content = `PhysicalVolume lvm_name }
160               | None ->
161                   disk (* Spare/unknown. *)
162   ) m_disks in
163
164   (* Now we have either detected partitions or a filesystem on each
165    * physical device (or perhaps neither).  See what is on those
166    * partitions.
167    *)
168   let m_disks = List.map (
169     function
170     | ({ d_dev = dev; d_content = `Partitions parts } as disk) ->
171         let ps = List.map (
172           fun p ->
173             if p.part_status = Bootable || p.part_status = Nonbootable then (
174               let fs = probe_for_filesystem p.part_dev in
175               match fs with
176               | Some fs ->
177                   { p with part_content = `Filesystem fs }
178               | None ->
179                   (* Is it a PV? *)
180                   let pv = probe_for_pv p.part_dev in
181                   match pv with
182                   | Some lvm_name ->
183                       { p with part_content = `PhysicalVolume lvm_name }
184                   | None ->
185                       p (* Spare/unknown. *)
186             ) else p
187         ) parts.parts in
188         let parts = { parts with parts = ps } in
189         { disk with d_content = `Partitions parts }
190     | disk -> disk
191   ) m_disks in
192
193   (* LVM filesystem detection
194    *
195    * Look for all disks/partitions which have been identified as PVs
196    * and pass those back to the respective LVM plugin for LV detection.
197    *
198    * (Note - a two-stage process because an LV can be spread over
199    * several PVs, so we have to detect all PVs belonging to a
200    * domain first).
201    *
202    * XXX To deal with RAID (ie. md devices) we will need to loop
203    * around here because RAID is like LVM except that they normally
204    * present as block devices which can be used by LVM.
205    *)
206   (* First: LV detection.
207    * Find all physical volumes, can be disks or partitions.
208    *)
209   let pvs_on_disks = List.filter_map (
210     function
211     | { d_dev = d_dev;
212         d_content = `PhysicalVolume pv } -> Some (pv, (d_dev :> device))
213     | _ -> None
214   ) m_disks in
215   let pvs_on_partitions = List.map (
216     function
217     | { d_content = `Partitions { parts = parts } } ->
218         List.filter_map (
219           function
220           | { part_dev = part_dev;
221               part_content = `PhysicalVolume pv } ->
222               Some (pv, part_dev)
223           | _ -> None
224         ) parts
225     | _ -> []
226   ) m_disks in
227   let lvs = List.concat (pvs_on_disks :: pvs_on_partitions) in
228
229   (* Second: filesystem on LV detection.
230    * Group the LVs by plug-in type.
231    *)
232   let cmp (a,_) (b,_) = compare a b in
233   let lvs = List.sort ~cmp lvs in
234   let lvs = group_by lvs in
235
236   let lvs =
237     List.map (fun (pv, devs) -> list_lvs pv.lvm_plugin_id devs) lvs in
238   let lvs = List.concat lvs in
239
240   (* lvs is a list of potential LV devices.  Now run them through the
241    * probes to see if any contain filesystems.
242    *)
243   let filesystems =
244     List.filter_map (
245       fun ({ lv_dev = dev } as lv) ->
246         match probe_for_filesystem dev with
247         | Some fs -> Some (lv, fs)
248         | None -> None
249     ) lvs in
250
251   { machine with
252       m_disks = m_disks;
253       m_lv_filesystems = filesystems }
254
255 (*----------------------------------------------------------------------*)
256
257 (* We describe the ownership of each part of the disk using a
258  * segment tree. http://en.wikipedia.org/wiki/Segment_tree
259  *
260  * Note that each part can (and usually is) owned multiple times
261  * (eg. by a filesystem and by the partition that the filesystem
262  * lies inside).  Also, the segment tree is effectively read-only.
263  * We build it up as a final step given the flat list of segments
264  * identified by the algorithm in 'iter_over_machine'.
265  *)
266
267 (* General binary tree type.  Data 'a is stored in the leaves and 'b
268  * is stored in the nodes.
269  *)
270 type ('a,'b) binary_tree =
271   | Leaf of 'a
272   | Node of ('a,'b) binary_tree * 'b * ('a,'b) binary_tree
273
274 (* This prints out the binary tree in graphviz dot format. *)
275 let print_binary_tree leaf_printer node_printer tree =
276   (* Assign a unique, fixed label to each node. *)
277   let label =
278     let i = ref 0 in
279     let hash = Hashtbl.create 13 in
280     fun node ->
281       try Hashtbl.find hash node
282       with Not_found ->
283         let i = incr i; !i in
284         let label = "n" ^ string_of_int i in
285         Hashtbl.add hash node label;
286         label
287   in
288   (* Recursively generate the graphviz file. *)
289   let rec print = function
290     | (Leaf a as leaf) ->
291         eprintf "  %s [shape=box, label=\"%s\"];\n"
292           (label leaf) (leaf_printer a)
293     | (Node (left,b,right) as node) ->
294         eprintf "  %s [label=\"%s\"];\n"
295           (label node) (node_printer b);
296         eprintf "  %s -> %s [tailport=sw];\n" (label node) (label left);
297         eprintf "  %s -> %s [tailport=se];\n" (label node) (label right);
298         print left;
299         print right;
300   in
301   eprintf "/* Use 'dot -Tpng foo.dot > foo.png' to convert to a png file. */\n";
302   eprintf "digraph G {\n";
303   print tree;
304   eprintf "}\n";
305
306 (* A segment describes the owner of a range of disk addresses. *)
307 type segment =
308     int63 * int63 *                     (* disk offset, size of segment *)
309       [ `Filesystem of filesystem
310       | `Partitions of partitions
311       | `PhysicalVolume of pv ] *       (* owner *)
312       int63                             (* owner offset *)
313
314 type interval = int63 * int63           (* start point, end point (bytes) *)
315
316 (* The special segment tree structure that we construct in create_ownership. *)
317 type segment_tree =
318     (interval * segment list, interval * segment list) binary_tree
319
320 type ownership =
321     (device *                           (* block_device (disk) *)
322        segment_tree) list               (* segment tree for this disk *)
323
324 (* List of owned segments before we build the segment tree. *)
325 type ownership_list =
326     (device *                           (* block_device (disk) *)
327        segment) list
328
329 (* Ownership tables. *)
330 let create_ownership machine =
331   (* Iterate over all the things which can claim ownership of a
332    * disk block (filesystems, partitions, PVs).
333    *)
334   let rec iter_over_machine
335       ({m_disks = disks; m_lv_filesystems = lv_filesystems} as machine) =
336
337     (* No segments to begin with. *)
338     let ownership = [] in
339
340     (* Iterate over disks. *)
341     let ownership =
342       List.fold_left (
343         fun ownership ->
344           function
345           | { d_content = (`Filesystem fs as owner) } ->
346               iter_over_filesystem machine ownership fs owner
347           | { d_content = (`Partitions parts as owner) } ->
348               iter_over_partitions machine ownership parts owner
349           | { d_content = (`PhysicalVolume pv as owner) } ->
350               iter_over_pv machine ownership pv owner
351           | { d_content = `Unknown } -> ownership
352       ) ownership disks in
353
354     (* Iterate over LV filesystems. *)
355     let ownership =
356       List.fold_left (
357         fun ownership (lv, fs) ->
358           let owner = `Filesystem fs in
359           iter_over_filesystem machine ownership fs owner
360       ) ownership lv_filesystems in
361
362     ownership
363
364   (* Iterate over the blocks in a single filesystem. *)
365   and iter_over_filesystem machine ownership {fs_dev = dev} owner =
366     iter_over_device machine ownership dev owner
367
368   (* Iterate over the blocks in a set of partitions, then
369    * iterate over the contents of the partitions.
370    *)
371   and iter_over_partitions machine ownership
372       {parts = parts; parts_dev = parts_dev} owner =
373     let ownership = iter_over_device machine ownership parts_dev owner in
374
375     let ownership =
376       List.fold_left (
377         fun ownership ->
378           function
379           | { part_content = (`Filesystem fs as owner) } ->
380               iter_over_filesystem machine ownership fs owner
381           | { part_content = (`PhysicalVolume pv as owner) } ->
382               iter_over_pv machine ownership pv owner
383           | { part_content = `Unknown } -> ownership
384       ) ownership parts in
385
386     ownership
387
388   (* Iterate over the blocks in a PV. *)
389   and iter_over_pv machine ownership {pv_dev = dev} owner =
390     iter_over_device machine ownership dev owner
391
392   (* Iterate over the blocks in a device, assigning ownership to 'owner'
393    *
394    * In reality (1): There can be several owners for each block, so we
395    * incrementally add ownership to the ownership_list (which eventually
396    * will be turned into a segment tree).
397    * In reality (2): Iterating over blocks would take ages and result
398    * in a very inefficient ownership representation.  Instead we look
399    * at minimum contiguous extents.
400    *)
401   and iter_over_device { m_disks = disks } ownership dev owner =
402     let size = dev#size in
403     let disks = List.map (fun {d_dev = dev} -> (dev :> device)) disks in
404
405     let rec loop ownership offset =
406       if offset < size then (
407         let devs, extent = get_next_extent disks dev offset in
408         if devs = [] then
409           eprintf "warning: no device found under %s\n"
410             (string_of_owner owner);
411         let ownership =
412           List.fold_left (
413             fun ownership (disk, disk_offset) ->
414               let elem = disk, (disk_offset, extent, owner, offset) in
415               elem :: ownership
416           ) ownership devs in
417         loop ownership (offset +^ extent)
418       )
419       else ownership
420     in
421     loop ownership ~^0
422
423   (* Return the length of the next contiguous region in the device starting
424    * at the given byte offset.  Also return the underlying block device(s)
425    * if there is one.
426    *)
427   and get_next_extent disks (dev : device) offset =
428     let this_extent = dev#contiguous offset in
429
430     (* If this disk is a block_device (a member of the 'disks' list)
431      * then we've hit the bottom layer of devices, so just return it.
432      *)
433     if List.memq dev disks then
434       [dev, offset], this_extent
435     else (
436       let blocksize = dev#blocksize in
437       let block = offset /^ blocksize in
438       let offset_in_block = offset -^ block *^ blocksize in
439
440       (* Map from this block to the devices one layer down. *)
441       let devs = dev#map_block block in
442
443       (* Get the real device offsets, adding the offset from start of block. *)
444       let devs =
445         List.map
446           (fun (dev, dev_offset) -> dev, dev_offset +^ offset_in_block)
447           devs in
448
449       let devs =
450         List.map
451           (fun (dev, dev_offset) ->
452              get_next_extent disks dev dev_offset)
453           devs in
454
455       (* Work out the minimum contiguous extent from this offset. *)
456       let devs, extent =
457         let extents = List.map snd devs in
458         let devs = List.concat (List.map fst devs) in
459         let extent = List.fold_left min this_extent extents in
460         devs, extent in
461
462       devs, extent
463     )
464
465   and string_of_owner = function
466     | `Filesystem {fs_plugin_id = fs_plugin_id; fs_dev = fs_dev} ->
467         sprintf "%s(%s)" fs_dev#name fs_plugin_id
468     | `PhysicalVolume { pv_uuid = pv_uuid } ->
469         "PV:" ^ pv_uuid
470     | `Partitions { parts_plugin_id = parts_plugin_id } ->
471         parts_plugin_id
472   in
473
474   (* Build the list of segments. *)
475   let ownership : ownership_list = iter_over_machine machine in
476
477   (* Group the segments together by disk. *)
478   let ownership =
479     let ownership = List.sort ownership in
480     group_by ownership in
481
482   (* If debugging, print the segments that we found. *)
483   if !debug then (
484     List.iter (
485       fun (disk, segments) ->
486         eprintf "ownership segment list of %s %s:\n" machine.m_name disk#name;
487         List.iter (
488           fun (disk_offset, size, owner, owner_offset) ->
489             let blocksize = disk#blocksize in
490             let disk_offset_in_blocks, disk_offset_in_block =
491               disk_offset /^ blocksize, disk_offset %^ blocksize in
492             let size_in_blocks, size_in_block =
493               size /^ blocksize, size %^ blocksize in
494
495             eprintf "  %s[%s:%s] %s[%s:%s] %s@%s\n"
496               (Int63.to_string disk_offset)
497                 (Int63.to_string disk_offset_in_blocks)
498                 (Int63.to_string disk_offset_in_block)
499               (Int63.to_string size)
500                 (Int63.to_string size_in_blocks)
501                 (Int63.to_string size_in_block)
502               (string_of_owner owner)
503               (Int63.to_string owner_offset)
504         ) segments
505     ) ownership
506   );
507
508   (* Build the segment tree from the ownership list (of segments).
509    * For an explanation of this process see:
510    * http://en.wikipedia.org/wiki/Segment_tree
511    *)
512   let ownership =
513     List.map (
514       fun (disk, segments) ->
515         (* Construct the list of distinct endpoints. *)
516         let eps =
517           List.map
518             (fun (start, size, _, _) -> [start; start +^ size])
519             segments in
520         let eps = sort_uniq (List.concat eps) in
521
522         (* Construct the elementary intervals. *)
523         let elints =
524           let elints, lastpoint =
525             List.fold_left (
526               fun (elints, prevpoint) point ->
527                 ((point, point) :: (prevpoint, point) :: elints), point
528             ) ([], Int63.min_int) eps in
529           let elints = (lastpoint, Int63.max_int) :: elints in
530           List.rev elints in
531
532         if !debug then (
533           eprintf "elementary intervals for %s (%d in total):\n"
534             disk#name (List.length elints);
535           List.iter (
536             fun (startpoint, endpoint) ->
537               eprintf "  %s %s\n"
538                 (Int63.to_string startpoint) (Int63.to_string endpoint)
539           ) elints
540         );
541
542         (* Construct the binary tree of elementary intervals. *)
543         let tree =
544           (* Each elementary interval becomes a leaf. *)
545           let elints = List.map (fun elint -> Leaf elint) elints in
546           (* Recursively build this into a binary tree. *)
547           let rec make_layer = function
548             | [] -> []
549             | ([_] as x) -> x
550             (* Turn pairs of leaves at the bottom level into nodes. *)
551             | (Leaf _ as a) :: (Leaf _ as b) :: xs ->
552                 let xs = make_layer xs in
553                 Node (a, (), b) :: xs
554             (* Turn pairs of nodes at higher levels into nodes. *)
555             | (Node _ as left) :: ((Node _|Leaf _) as right) :: xs ->
556                 let xs = make_layer xs in
557                 Node (left, (), right) :: xs
558             | Leaf _ :: _ -> assert false (* never happens??? (I think) *)
559           in
560           let rec loop = function
561             | [] -> assert false
562             | [x] -> x
563             | xs -> loop (make_layer xs)
564           in
565           loop elints in
566
567         if !debug then (
568           let leaf_printer (startpoint, endpoint) =
569             sprintf "%s-%s"
570               (Int63.to_string startpoint) (Int63.to_string endpoint)
571           in
572           let node_printer () = "" in
573           print_binary_tree leaf_printer node_printer tree
574         );
575
576         (* Insert the segments into the tree one by one. *)
577         let tree =
578           (* For each node/leaf in the tree, add its interval and an
579            * empty list which will be used to store the segments.
580            *)
581           let rec interval_tree = function
582             | Leaf elint -> Leaf (elint, [])
583             | Node (left, (), right) ->
584                 let left = interval_tree left in
585                 let right = interval_tree right in
586                 let (leftstart, _) = interval_of_node left in
587                 let (_, rightend) = interval_of_node right in
588                 let interval = leftstart, rightend in
589                 Node (left, (interval, []), right)
590           and interval_of_node = function
591             | Leaf (elint, _) -> elint
592             | Node (_, (interval, _), _) -> interval
593           in
594
595           let tree = interval_tree tree in
596           (* This should always be true: *)
597           assert (interval_of_node tree = (Int63.min_int, Int63.max_int));
598
599           (* "Contained in" operator.
600            * 'a <-< b' iff 'a' is a subinterval of 'b'.
601            *      |<---- a ---->|
602            * |<----------- b ----------->|
603            *)
604           let (<-<) (a1, a2) (b1, b2) = b1 <= a1 && a2 <= b2 in
605
606           (* "Intersects" operator.
607            * 'a /\ b' iff intervals 'a' and 'b' overlap, eg:
608            *      |<---- a ---->|
609            *                |<----------- b ----------->|
610            *)
611           let ( /\ ) (a1, a2) (b1, b2) = a2 > b1 || b2 > a1 in
612
613           let rec insert_segment tree segment =
614             let start, size, owner, owner_offset = segment in
615             let seginterval = start, start +^ size in
616             eprintf "inserting (%s,%s) ...\n"
617               (Int63.to_string (fst seginterval))
618               (Int63.to_string (snd seginterval));
619             match tree with
620             (* Test if we should insert into this leaf or node: *)
621             | Leaf (interval, segs) when interval <-< seginterval ->
622                 Leaf (interval, segment :: segs)
623             | Node (left, (interval, segs), right)
624                 when interval <-< seginterval ->
625                 Node (left, (interval, segment :: segs), right)
626
627             | (Leaf _) as leaf -> leaf
628
629             (* Else, should we insert into left or right subtrees? *)
630             | Node (left, i, right) ->
631                 let left =
632                   if seginterval /\ interval_of_node left then
633                     insert_segment left segment
634                   else
635                     left in
636                 let right =
637                   if seginterval /\ interval_of_node right then
638                     insert_segment right segment
639                   else
640                     right in
641                 Node (left, i, right)
642           in
643           let tree = List.fold_left insert_segment tree segments in
644           tree in
645
646         if !debug then (
647           let printer ((sp, ep), segments) =
648             sprintf "[%s-%s] " (Int63.to_string sp) (Int63.to_string ep) ^
649               String.concat ";"
650               (List.map (fun (_, _, owner,_) -> string_of_owner owner)
651                  segments)
652           in
653           print_binary_tree printer printer tree
654         );
655         (disk, tree)
656     ) ownership in
657
658   (* Return the ownership structure. *)
659   ownership