Automatically sets it to the smaller of:
* Size of the swap partition on source machine
* 10% of instance disk

Signed-off-by: Ben Lipton <[email protected]>
---
 p2v-transfer/p2v_transfer.py           |   59 ++++++++++++++-----
 p2v-transfer/test/p2v_transfer_test.py |   99 ++++++++++++++++++++++++++++----
 2 files changed, 133 insertions(+), 25 deletions(-)

diff --git a/p2v-transfer/p2v_transfer.py b/p2v-transfer/p2v_transfer.py
index b2fc12f..432051f 100755
--- a/p2v-transfer/p2v_transfer.py
+++ b/p2v-transfer/p2v_transfer.py
@@ -192,34 +192,63 @@ def _WaitForCompletion(channel):
     time.sleep(.01)
 
 
-def PartitionTargetDisks(client, swap_cyls):
-  """Partition and format the disks on the target machine.
+def GetDiskSize(client):
+  """Determine how much disk is available, how much swap space to include.
 
-  Sends commands over the SSH connection to partition and format the
-  disk of the target instance.
+  For swap size, returns the minimum of:
+  - amount of swap space on the source machine
+  - 10% of the target drive
 
   @type client: paramiko.SSHClient
   @param client: SSH client object used to connect to the instance.
-  @type swap_cyls: int
-  @param swap_cyls: Desired size of swap space, in cylinders
+  @rtype: (int, int)
+  @return: Total size in megabytes, swap size in megabytes
 
   """
-  # Find out how many cylinders are available on target
-  total_cyls = 0
-  stdin, stdout, stderr = client.exec_command("sfdisk -l /dev/xvda")
+   # Find out how many MB are available on target
+  stdin, stdout, stderr = client.exec_command("fdisk -l /dev/xvda")
   for line in stdout:
     if line.startswith("Disk /dev/xvda:"):
       words = line.split()
-      total_cyls = int(words[2])
+      total_megs = int(words[2])
       break
   stdout.close()
 
-  nonswap_cyls = total_cyls - swap_cyls
-  sfdisk_command = """sfdisk /dev/xvda <<EOF
+  swap_output = subprocess.Popen(["swapon", "-s"],
+                                 stdout=subprocess.PIPE,
+                                 stderr=subprocess.PIPE).communicate()[0]
+  swap_megs = 0
+  for line in swap_output.splitlines()[1:]:
+    words = line.split()
+    try:
+      size_kb = int(words[2])
+      swap_megs += size_kb / 1024
+    except ValueError:
+      pass  # This line doesn't have data on it
+
+  return total_megs, min(swap_megs, int(total_megs * 0.1))
+
+
+def PartitionTargetDisks(client, total_megs, swap_megs):
+  """Partition and format the disks on the target machine.
+
+  Sends commands over the SSH connection to partition and format the
+  disk of the target instance.
+
+  @type client: paramiko.SSHClient
+  @param client: SSH client object used to connect to the instance.
+  @type total_megs: int
+  @param total_megs: Total size of disk, in megabytes
+  @type swap_megs: int
+  @param swap_megs: Desired size of swap space, in megabytes
+
+  """
+  nonswap_megs = total_megs - swap_megs
+  sfdisk_command = """sfdisk -uM /dev/xvda <<EOF
 0,%d,83
 ,,82
 EOF
-""" % nonswap_cyls
+""" % nonswap_megs
 
   _RunCommandAndWait(client, sfdisk_command)
 
@@ -232,6 +261,7 @@ EOF
 
   _RunCommandAndWait(client, " && ".join(other_commands))
 
+
 def TransferFiles(user, host, keyfile):
   """Transfer files to the bootstrap OS.
 
@@ -295,7 +325,8 @@ def main(argv):
     client = EstablishConnection(user, host, key)
     MountSourceFilesystems(root_dev)
     if VerifyKernelMatches(client):
-      PartitionTargetDisks(client, 10)
+      total_megs, swap_megs = GetDiskSize(client)
+      PartitionTargetDisks(client, total_megs, swap_megs)
       TransferFiles(user, host, keyfile)
       RunFixScripts(client)
       ShutDownTarget(client)
diff --git a/p2v-transfer/test/p2v_transfer_test.py 
b/p2v-transfer/test/p2v_transfer_test.py
index ebfb1e1..c3d580a 100755
--- a/p2v-transfer/test/p2v_transfer_test.py
+++ b/p2v-transfer/test/p2v_transfer_test.py
@@ -61,6 +61,8 @@ class P2vtransferTest(unittest.TestCase):
       self.host,
       self.pkeyfile,
       ]
+    self.swapsize = 1024
+    self.totsize = 102400
 
   def _MockRunCommandAndWait(self, command, exit_status=0):
     stdin = _MockChannelFile(self.mox)
@@ -86,6 +88,7 @@ class P2vtransferTest(unittest.TestCase):
   def _StubOutAllModuleFunctions(self):
     self.module_functions = [
       "EstablishConnection",
+      "GetDiskSize",
       "PartitionTargetDisks",
       "MountSourceFilesystems",
       "TransferFiles",
@@ -108,22 +111,91 @@ class P2vtransferTest(unittest.TestCase):
     self.module.ShutDownTarget(self.client)
     self.mox.VerifyAll()
 
-  def testPartitionTargetDisksSendsCommands(self):
-    swap_cyls = 10
-    tot_cyls = 600
+  def testGetDiskSizeHandlesLargeDisk(self):
+    """On large disks, swap size should be the same as the source.
+
+    uses self.totsize and self.swapsize to generate command outputs such that
+    the returned values will be approximately self.totsize and self.swapsize.
+    """
+    popen = self.mox.CreateMock(self.module.subprocess.Popen)
+    self.mox.StubOutWithMock(self.module.subprocess, 'Popen',
+                             use_mock_anything=True)
+
+    tot_bytes = self.totsize * 1024 * 1024
+    swap_kb = self.swapsize * 1024
+
+    fdisk_output = "Disk /dev/xvda: %s MB, %d bytes" % (self.totsize,
+                                                        tot_bytes)
+    swapon_output = """
+Filename                       Type            Size    Used    Priority
+/dev/sda5                       partition      %d      0       -1
+""" % swap_kb
 
-    sfdisk_output = "Disk /dev/xvda: %d cylinders, 255 heads, etc." % tot_cyls
     stdout = _MockChannelFile(self.mox)
-    stdout._SetOutput(sfdisk_output)
+    stdout._SetOutput(fdisk_output)
+
+    self.client.exec_command("fdisk -l /dev/xvda").AndReturn((None, stdout,
+                                                              None))
+    call = self.module.subprocess.Popen(["swapon", "-s"],
+                                        stdout=self.module.subprocess.PIPE,
+                                        stderr=self.module.subprocess.PIPE)
+    call.AndReturn(popen)
+    popen.communicate().AndReturn((swapon_output, None))
+
+    self.mox.ReplayAll()
+    total, swap = self.module.GetDiskSize(self.client)
+    self.assertEqual(total, self.totsize)
+    # Should return same swap size as source machine
+    self.assertEqual(swap, self.swapsize)
+    self.mox.VerifyAll()
+
+  def testGetDiskSizeHandlesSmallDisk(self):
+    """On small disks, swap size should be 10% of the disk.
+
+    Makes self.totsize small, so that the returned swap size will be smaller
+    than self.swapsize.
+    """
+    self.totsize = self.swapsize * 5
 
-    self.client.exec_command("sfdisk -l /dev/xvda").AndReturn((None, stdout,
+    popen = self.mox.CreateMock(self.module.subprocess.Popen)
+    self.mox.StubOutWithMock(self.module.subprocess, 'Popen',
+                             use_mock_anything=True)
+
+    tot_bytes = self.totsize * 1024 * 1024
+    swap_kb = self.swapsize * 1024
+
+    fdisk_output = "Disk /dev/xvda: %s MB, %d bytes" % (self.totsize,
+                                                         tot_bytes)
+    swapon_output = """
+Filename                       Type            Size    Used    Priority
+/dev/sda5                       partition      %d      0       -1
+""" % swap_kb
+
+    stdout = _MockChannelFile(self.mox)
+    stdout._SetOutput(fdisk_output)
+
+    self.client.exec_command("fdisk -l /dev/xvda").AndReturn((None, stdout,
                                                                None))
+    call = self.module.subprocess.Popen(["swapon", "-s"],
+                                        stdout=self.module.subprocess.PIPE,
+                                        stderr=self.module.subprocess.PIPE)
+    call.AndReturn(popen)
+    popen.communicate().AndReturn((swapon_output, None))
 
-    sfdisk_command = """sfdisk /dev/xvda <<EOF
+    self.mox.ReplayAll()
+    total, swap = self.module.GetDiskSize(self.client)
+    self.assertEqual(total, self.totsize)
+    # swap size should be about 10% of the total
+    self.assertEqual(swap, self.totsize/10)
+    self.assertNotEqual(swap, self.swapsize)
+    self.mox.VerifyAll()
+
+  def testPartitionTargetDisksSendsCommands(self):
+    sfdisk_command = """sfdisk -uM /dev/xvda <<EOF
 0,%d,83
 ,,82
 EOF
-""" % (tot_cyls - swap_cyls)
+""" % (self.totsize - self.swapsize)
 
     self._MockRunCommandAndWait(sfdisk_command)
 
@@ -135,7 +207,7 @@ EOF
     self._MockRunCommandAndWait(commands)
 
     self.mox.ReplayAll()
-    self.module.PartitionTargetDisks(self.client, swap_cyls)
+    self.module.PartitionTargetDisks(self.client, self.totsize, self.swapsize)
     self.mox.VerifyAll()
 
   def testTransferFilesExitsOnError(self):
@@ -187,7 +259,9 @@ EOF
                                     self.pkey).AndReturn(self.client)
     self.module.MountSourceFilesystems(self.root_dev)
     self.module.VerifyKernelMatches(self.client).AndReturn(True)
-    self.module.PartitionTargetDisks(self.client, 10)
+    self.module.GetDiskSize(self.client).AndReturn((self.totsize,
+                                                    self.swapsize))
+    self.module.PartitionTargetDisks(self.client, self.totsize, self.swapsize)
     self.module.TransferFiles("root", self.host, self.pkeyfile)
     self.module.RunFixScripts(self.client)
     self.module.ShutDownTarget(self.client)
@@ -228,7 +302,10 @@ EOF
                                     self.pkey).AndReturn(self.client)
     self.module.MountSourceFilesystems(self.root_dev)
     self.module.VerifyKernelMatches(self.client).AndReturn(True)
-    call = self.module.PartitionTargetDisks(self.client, 10)
+    self.module.GetDiskSize(self.client).AndReturn((self.totsize,
+                                                    self.swapsize))
+    call = self.module.PartitionTargetDisks(self.client, self.totsize,
+                                            self.swapsize)
     call.AndRaise(self.module.P2VError("meep"))
     # Transfer is cancelled because of the error, but still we have:
     self.module.UnmountSourceFilesystems()
-- 
1.7.3.1

Reply via email to