Add support for snapshot size calculations

- govc: Add size (`s`) flag to `snapshot.tree` subcommand
- Add `object.SnapshotSize` and helper functions
  to support calculating snapshot size from other
  client code

Credit to @dougm for providing successive code examples,
explanation and guidance for understanding the required
steps to implement this support.

As noted on vmwareGH-2243, code provided by @dougm was (AFAIK)
based heavily on existing C# code for PowerCLI
implementation of the `Get-Snapshot` cmdlet.

- credit: @dougm
- refs vmware#2243

govc/vm/snapshot/tree.go

@@ -1,5 +1,5 @@
 /*
-Copyright (c) 2016 VMware, Inc. All Rights Reserved.
+Copyright (c) 2021 VMware, Inc. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -25,6 +25,8 @@ import (
 
 	"github.com/vmware/govmomi/govc/cli"
 	"github.com/vmware/govmomi/govc/flags"
+	"github.com/vmware/govmomi/object"
+	"github.com/vmware/govmomi/units"
 	"github.com/vmware/govmomi/vim25/mo"
 	"github.com/vmware/govmomi/vim25/types"
 )
@@ -38,8 +40,10 @@ type tree struct {
 	description bool
 	fullPath    bool
 	id          bool
+	size        bool
 
-	info *types.VirtualMachineSnapshotInfo
+	info   *types.VirtualMachineSnapshotInfo
+	layout *types.VirtualMachineFileLayoutEx
 }
 
 func init() {
@@ -59,6 +63,7 @@ func (cmd *tree) Register(ctx context.Context, f *flag.FlagSet) {
 	f.BoolVar(&cmd.fullPath, "f", false,
 		"Print the full path prefix for snapshot")
 	f.BoolVar(&cmd.id, "i", false, "Print the snapshot id")
+	f.BoolVar(&cmd.size, "s", false, "Print the snapshot size")
 }
 
 func (cmd *tree) Description() string {
@@ -78,7 +83,7 @@ func (cmd *tree) Process(ctx context.Context) error {
 	return nil
 }
 
-func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapshotTree) {
+func (cmd *tree) write(level int, parent string, pref *types.ManagedObjectReference, st []types.VirtualMachineSnapshotTree) {
 	for _, s := range st {
 		sname := s.Name
 
@@ -92,7 +97,10 @@ func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapsh
 			names = append(names, sname)
 		}
 
+		isCurrent := false
+
 		if s.Snapshot == *cmd.info.CurrentSnapshot {
+			isCurrent = true
 			if cmd.current {
 				names = append(names, ".")
 			} else if cmd.currentName {
@@ -105,6 +113,12 @@ func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapsh
 			var attr []string
 			var meta string
 
+			if cmd.size {
+				size := object.SnapshotSize(s.Snapshot, pref, cmd.layout, isCurrent)
+
+				attr = append(attr, units.ByteSize(size).String())
+			}
+
 			if cmd.id {
 				attr = append(attr, s.Snapshot.Value)
 			}
@@ -127,7 +141,7 @@ func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapsh
 			}
 		}
 
-		cmd.write(level+2, sname, s.ChildSnapshotList)
+		cmd.write(level+2, sname, &s.Snapshot, s.ChildSnapshotList)
 	}
 }
 
@@ -147,7 +161,7 @@ func (cmd *tree) Run(ctx context.Context, f *flag.FlagSet) error {
 
 	var o mo.VirtualMachine
 
-	err = vm.Properties(ctx, vm.Reference(), []string{"snapshot"}, &o)
+	err = vm.Properties(ctx, vm.Reference(), []string{"snapshot", "layoutEx"}, &o)
 	if err != nil {
 		return err
 	}
@@ -161,8 +175,9 @@ func (cmd *tree) Run(ctx context.Context, f *flag.FlagSet) error {
 	}
 
 	cmd.info = o.Snapshot
+	cmd.layout = o.LayoutEx
 
-	cmd.write(0, "", o.Snapshot.RootSnapshotList)
+	cmd.write(0, "", nil, o.Snapshot.RootSnapshotList)
 
 	return nil
 }

object/virtual_machine.go

@@ -1,5 +1,5 @@
 /*
-Copyright (c) 2015-2017 VMware, Inc. All Rights Reserved.
+Copyright (c) 2015-2021 VMware, Inc. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -40,6 +40,34 @@ type VirtualMachine struct {
 	Common
 }
 
+// extractDiskLayoutFiles is a helper function used to extract file keys for
+// all disk files attached to the virtual machine at the current point of
+// running.
+func extractDiskLayoutFiles(diskLayoutList []types.VirtualMachineFileLayoutExDiskLayout) []int {
+	var result []int
+
+	for _, layoutExDisk := range diskLayoutList {
+		for _, link := range layoutExDisk.Chain {
+			for i := range link.FileKey { // diskDescriptor, diskExtent pairs
+				result = append(result, int(link.FileKey[i]))
+			}
+		}
+	}
+
+	return result
+}
+
+// removeKey is a helper function for removing a specific file key from a list
+// of keys associated with disks attached to a virtual machine.
+func removeKey(l *[]int, key int) {
+	for i, k := range *l {
+		if k == key {
+			*l = append((*l)[:i], (*l)[i+1:]...)
+			break
+		}
+	}
+}
+
 func NewVirtualMachine(c *vim25.Client, ref types.ManagedObjectReference) *VirtualMachine {
 	return &VirtualMachine{
 		Common: NewCommon(c, ref),
@@ -628,6 +656,63 @@ func (m snapshotMap) add(parent string, tree []types.VirtualMachineSnapshotTree)
 	}
 }
 
+// SnapshotSize calculates the size of a given snapshot in bytes. If the
+// snapshot is current, disk files not associated with any parent snapshot are
+// included in size calculations. This allows for measuring and including the
+// growth from the last fixed snapshot to the present state.
+func SnapshotSize(info types.ManagedObjectReference, parent *types.ManagedObjectReference, vmlayout *types.VirtualMachineFileLayoutEx, isCurrent bool) int {
+	var fileKeyList []int
+	var parentFiles []int
+	var allSnapshotFiles []int
+
+	diskFiles := extractDiskLayoutFiles(vmlayout.Disk)
+
+	for _, layout := range vmlayout.Snapshot {
+		diskLayout := extractDiskLayoutFiles(layout.Disk)
+		allSnapshotFiles = append(allSnapshotFiles, diskLayout...)
+
+		if layout.Key.Value == info.Value {
+			fileKeyList = append(fileKeyList, int(layout.DataKey)) // The .vmsn file
+			fileKeyList = append(fileKeyList, diskLayout...)       // The .vmdk files
+		} else if parent != nil && layout.Key.Value == parent.Value {
+			parentFiles = append(parentFiles, diskLayout...)
+		}
+	}
+
+	for _, parentFile := range parentFiles {
+		removeKey(&fileKeyList, parentFile)
+	}
+
+	for _, file := range allSnapshotFiles {
+		removeKey(&diskFiles, file)
+	}
+
+	fileKeyMap := make(map[int]types.VirtualMachineFileLayoutExFileInfo)
+	for _, file := range vmlayout.File {
+		fileKeyMap[int(file.Key)] = file
+	}
+
+	size := 0
+
+	for _, fileKey := range fileKeyList {
+		file := fileKeyMap[fileKey]
+		if parent != nil ||
+			(file.Type != string(types.VirtualMachineFileLayoutExFileTypeDiskDescriptor) &&
+				file.Type != string(types.VirtualMachineFileLayoutExFileTypeDiskExtent)) {
+			size += int(file.Size)
+		}
+	}
+
+	if isCurrent {
+		for _, diskFile := range diskFiles {
+			file := fileKeyMap[diskFile]
+			size += int(file.Size)
+		}
+	}
+
+	return size
+}
+
 // FindSnapshot supports snapshot lookup by name, where name can be:
 // 1) snapshot ManagedObjectReference.Value (unique)
 // 2) snapshot name (may not be unique)