1#!/usr/bin/env python3

2# Copyright 2019 Canonical Ltd.

3# All rights reserved.

4#

5# Written by:

6# Jonathan Cave <jonathan.cave@canonical.com>

7#

8# Perform bonnie++ disk test

9

10 from collections import namedtuple

11 from contextlib import ExitStack

12 import os

13 import subprocess as sp

14 import sys

15 import tempfile

16

17 import psutil

18

19

20 def mountpoint(device):

21 for part in psutil.disk_partitions():

22 if part.device == device:

23 return part.mountpoint

24 return None

25

26

27 def find_largest_partition(device):

28 BlkDev = namedtuple("BlkDev", ["name", "size", "type", "fstype"])

29 cmd = "lsblk -b -l -n -o NAME,SIZE,TYPE,FSTYPE {}".format(device)

30 out = sp.check_output(cmd, shell=True)

31 blk_devs = []

32 for entry in out.decode(sys.stdout.encoding).splitlines():

33 params = entry.strip().split()

34 if len(params) == 3:

35 # filesystem info missing, so it's unknown - skip

36 continue

37 blk_devs.append(BlkDev(*params))

38 blk_devs[:] = [

39 bd

40 for bd in blk_devs

41 if (bd.type in ("part", "md") and bd.fstype != "crypto_LUKS")

42 ]

43 if not blk_devs:

44 raise SystemExit(

45 "ERROR: No suitable partitions found on device {}".format(device)

46 )

47 blk_devs.sort(key=lambda bd: int(bd.size))

48 return blk_devs[-1].name

49

50

51 def mount(source, target):

52 cmd = "mount {} {}".format(source, target)

53 print("+", cmd, flush=True)

54 sp.check_call(cmd, shell=True)

55

56

57 def unmount(target):

58 cmd = "umount {}".format(target)

59 print("+", cmd, flush=True)

60 sp.check_call(cmd, shell=True)

61

62

63 def memory():

64 return psutil.virtual_memory().total / (1024 * 1024)

65

66

67 def free_space(test_dir):

68 du = psutil.disk_usage(test_dir)

69 return du.free / (1024 * 1024)

70

71

72 def devmapper_name(udev_name):

73 dm_name = None

74 sys_d = "/sys/block/{}".format(udev_name)

75 if os.path.isdir(os.path.join(sys_d, "dm")):

76 with open("/sys/block/{}/dm/name".format(udev_name), "r") as f:

77 dm_name = f.read().strip()

78 return dm_name

79

80

81 def run_bonnie(test_dir, user="root"):

82 # Set a maximum size on the amount of RAM, this has the effect of keeping

83 # the amount of data written during tests lower than default. This keeps

84 # duration of tests at something reasonable

85 force_mem_mb = 8000

86 if memory() < force_mem_mb: 86 ↛ 87 line 86 didn't jump to line 87 because the condition on line 86 was never true

87 force_mem_mb = memory()

88 # When running on disks with small drives (SSD/flash) we need to do

89 # some tweaking. Bonnie uses 2x RAM by default to write data. If that's

90 # more than available disk space, the test will fail inappropriately.

91 free = free_space(test_dir)

92 print("{}MB of free space available".format(free))

93 if (force_mem_mb * 2) > free: 93 ↛ 94 line 93 didn't jump to line 94 because the condition on line 93 was never true

94 force_mem_mb = free / 4

95 print("Forcing memory setting to {}MB".format(force_mem_mb))

96 cmd = "bonnie++ -d {} -u {} -r {}".format(test_dir, user, force_mem_mb)

97 print("+", cmd, flush=True)

98 sp.check_call(cmd, shell=True)

99

100

101 def devmapper_test(udev_name):

102 print("identified as a devmapper device...")

103 device = "/dev/{}".format(udev_name)

104 mount_dir = mountpoint(device)

105 if mount_dir:

106 print("{} already mounted at {}".format(device, mount_dir))

107 else:

108 dm_name = devmapper_name(udev_name)

109 if dm_name:

110 dm_device = os.path.join("/dev/mapper", dm_name)

111 if os.path.exists(dm_device):

112 mount_dir = mountpoint(dm_device)

113 if mount_dir:

114 print(

115 "{} already mounted at {}".format(dm_device, mount_dir)

116 )

117 with ExitStack() as stack:

118 if mount_dir is None:

119 mount_dir = tempfile.mkdtemp()

120 stack.callback(os.rmdir, mount_dir)

121 mount(device, mount_dir)

122 print("Performed mount of {} at {}".format(device, mount_dir))

123 stack.callback(unmount, mount_dir)

124 run_bonnie(mount_dir)

125

126

127 def disk_test(udev_name):

128 print("identified as a disk...")

129 device = "/dev/{}".format(udev_name)

130 part_to_test = "/dev/{}".format(find_largest_partition(device))

131 print("test will be run on partition {}".format(part_to_test))

132 mount_dir = mountpoint(part_to_test)

133 if mount_dir:

134 print("{} already mounted at {}".format(part_to_test, mount_dir))

135 with ExitStack() as stack:

136 if mount_dir is None:

137 mount_dir = tempfile.mkdtemp()

138 stack.callback(os.rmdir, mount_dir)

139 mount(part_to_test, mount_dir)

140 print("Performed mount {} at {}".format(part_to_test, mount_dir))

141 stack.callback(unmount, mount_dir)

142 run_bonnie(mount_dir)

143

144

145 def main():

146 udev_name = sys.argv[1]

147 print("Testing device {}".format(udev_name))

148

149 # Handle devmapper, and regular disks separately, and ignore mtdblock.

150 if udev_name.startswith("mtdblock"):

151 print("Ignoring mtdblock device")

152 elif devmapper_name(udev_name):

153 devmapper_test(udev_name)

154 else:

155 disk_test(udev_name)

156

157

158 if __name__ == "__main__":

159 main()