1#!/usr/bin/env python3

2

3 import os

4 import sys

5 import re

6

7 from argparse import ArgumentParser

8 from subprocess import Popen, PIPE

9

10

11 class MemoryTest:

12

13 def __init__(self):

14 self.free_memory = 0

15 self.system_memory = 0

16 self.swap_memory = 0

17 self.process_memory = 0

18 self.is_process_limited = False

19

20 @property

21 def threaded_memtest_script(self):

22 directory = os.path.dirname(os.path.abspath(__file__))

23 return os.path.join(directory, "threaded_memtest")

24

25 def _get_memory(self):

26 mem_info = open("/proc/meminfo", "r")

27 try:

28 while True:

29 line = mem_info.readline()

30 if line:

31 tokens = line.split()

32 if len(tokens) == 3:

33 if "MemTotal:" == tokens[0].strip():

34 self.system_memory = int(tokens[1].strip()) // 1024

35 elif tokens[0].strip() in [

36 "MemFree:",

37 "Cached:",

38 "Buffers:",

39 ]:

40 self.free_memory += int(tokens[1].strip()) // 1024

41 elif "SwapTotal:" == tokens[0].strip():

42 self.swap_memory = int(tokens[1].strip()) // 1024

43 else:

44 break

45 except Exception as e:

46 print(

47 "ERROR: Unable to get data from /proc/meminfo", file=sys.stderr

48 )

49 print(e, file=sys.stderr)

50 finally:

51 mem_info.close()

52

53 def _command(self, command, shell=True):

54 proc = Popen(command, shell=shell, stdout=PIPE, stderr=PIPE)

55 return proc

56

57 def _command_out(self, command, shell=True):

58 proc = self._command(command, shell)

59 return proc.communicate()[0].strip()

60

61 def get_limits(self):

62 self._get_memory()

63 print("System Memory: %u MB" % self.system_memory)

64 print("Free Memory: %u MB" % self.free_memory)

65 print("Swap Memory: %u MB" % self.swap_memory)

66

67 if self.system_memory == 0:

68 print("ERROR: could not determine system RAM", file=sys.stderr)

69 return False

70

71 # Process Memory

72 self.process_memory = self.free_memory

73 try:

74 arch = self._command_out("arch").decode()

75 if (

76 re.match(r"(i[0-9]86|s390|arm.*)", arch)

77 and self.free_memory > 1024

78 ):

79 self.is_process_limited = True

80 self.process_memory = 1024 # MB, due to 32-bit address space

81 print(

82 "%s arch, Limiting Process Memory: %u"

83 % (arch, self.process_memory)

84 )

85 # others? what about PAE kernel?

86 except Exception as e:

87 print(

88 "ERROR: could not determine system architecture via arch",

89 file=sys.stderr,

90 )

91 print(e, file=sys.stderr)

92 return False

93 return True

94

95 def run(self):

96 PASSED = 0

97 FAILED = 1

98

99 limits = self.get_limits()

100 if not limits:

101 return FAILED

102

103 # if process memory is limited, run multiple processes

104 if self.is_process_limited:

105 print("Running Multiple Process Memory Test")

106 if not self.run_multiple_process_test():

107 return FAILED

108 else:

109 print("Running Single Process Memory Test")

110 if not self.run_single_process_test():

111 return FAILED

112

113 # otherwised, passed

114 return PASSED

115

116 def run_single_process_test(self):

117 if not self.run_threaded_memory_test():

118 return False

119 return True

120

121 def run_multiple_process_test(self):

122 processes = self.free_memory // self.process_memory

123 # if not swap-less, add a process to hit swap

124 if not self.swap_memory == 0:

125 processes += 1

126 # check to make sure there's enough swap

127 required_memory = self.process_memory * processes

128 if required_memory > self.system_memory + self.swap_memory:

129 print(

130 "ERROR: this test requires a minimum of %u KB of swap "

131 "memory (%u configured)"

132 % (required_memory - self.system_memory, self.swap_memory),

133 file=sys.stderr,

134 )

135 print("Testing memory with %u processes" % processes)

136

137 print("Running threaded memory test:")

138 run_time = 60 # sec.

139 if not self.run_processes(

140 processes,

141 "%s -qv -m%um -t%u"

142 % (self.threaded_memtest_script, self.process_memory, run_time),

143 ):

144 print(

145 "Multi-process, threaded memory Test FAILED", file=sys.stderr

146 )

147 return False

148

149 return True

150

151 def run_threaded_memory_test(self):

152 # single-process threaded test

153 print("Starting Threaded Memory Test")

154

155 # run for Free Memory plus the lessor of 5% or 1GB

156 memory = (self.free_memory * 5) / 100

157 if memory > 1024: # MB

158 memory = 1024 # MB

159 memory = memory + self.free_memory

160 print("Running for %d MB total memory" % memory)

161

162 # run a test that will swap

163 if not self.swap_memory == 0:

164

165 # is there enough swap memory for the test?

166 if memory > self.system_memory + self.swap_memory:

167 print(

168 "ERROR: this test requires a minimum of %u KB of swap "

169 "memory (%u configured)"

170 % (memory - self.system_memory, self.swap_memory),

171 file=sys.stderr,

172 )

173 return False

174

175 # otherwise

176 run_time = 60 # sec.

177 print(

178 "Running for more than free memory at %u MB for %u sec."

179 % (memory, run_time)

180 )

181

182 command = "%s -qv -m%um -t%u" % (

183 self.threaded_memtest_script,

184 memory,

185 run_time,

186 )

187 print("Command is: %s" % command)

188 process = self._command(command)

189 process.communicate()

190 if process.returncode != 0:

191 print(

192 "%s returned code %s"

193 % (self.threaded_memtest_script, process.returncode),

194 file=sys.stderr,

195 )

196 print("More Than Free Memory Test failed", file=sys.stderr)

197 return False

198 print("More than free memory test complete.")

199

200 # run again for 15 minutes

201 print("Running for free memory")

202 process = self._command("%s -qv" % self.threaded_memtest_script)

203 process.communicate()

204 if process.returncode != 0:

205 print("Free Memory Test failed", file=sys.stderr)

206 else:

207 print("Free Memory Test succeeded")

208 sys.stdout.flush()

209 return process.returncode == 0

210

211 def run_processes(self, number, command):

212 passed = True

213 pipe = []

214 for i in range(number):

215 pipe.append(self._command(command))

216 print("Started: process %u pid %u: %s" % (i, pipe[i].pid, command))

217 sys.stdout.flush()

218 waiting = True

219 while waiting:

220 waiting = False

221 for i in range(number):

222 if pipe[i]:

223 line = pipe[i].communicate()[0]

224 if line and len(line) > 1:

225 print("process %u pid %u: %s" % (i, pipe[i].pid, line))

226 sys.stdout.flush()

227 if pipe[i].poll() == -1:

228 waiting = True

229 else:

230 return_value = pipe[i].poll()

231 if return_value != 0:

232 print(

233 "ERROR: process %u pid %u retuned %u"

234 % (i, pipe[i].pid, return_value),

235 file=sys.stderr,

236 )

237 passed = False

238 print(

239 "process %u pid %u returned success"

240 % (i, pipe[i].pid)

241 )

242 pipe[i] = None

243 sys.stdout.flush()

244 return passed

245

246

247 def main(args):

248 parser = ArgumentParser()

249 parser.add_argument(

250 "-q", "--quiet", action="store_true", help="Suppress output."

251 )

252 args = parser.parse_args(args)

253

254 if args.quiet:

255 sys.stdout = open(os.devnull, "a")

256 sys.stderr = open(os.devnull, "a")

257

258 test = MemoryTest()

259 return test.run()

260

261

262 if __name__ == "__main__":

263 sys.exit(main(sys.argv[1:]))