services/dockercleaner/tests/test_cleaner.py

   1 #!/usr/bin/env python3
   2 # Copyright (C) The Arvados Authors. All rights reserved.
   3 #
   4 # SPDX-License-Identifier: AGPL-3.0
   5
   6 import collections
   7 import itertools
   8 import json
   9 import os
  10 import random
  11 import tempfile
  12 import time
  13 import unittest
  14
  15 import docker
  16 from unittest import mock
  17
  18 from arvados_docker import cleaner
  19
  20 MAX_DOCKER_ID = (16 ** 64) - 1
  21
  22
  23 def MockDockerId():
  24     return '{:064x}'.format(random.randint(0, MAX_DOCKER_ID))
  25
  26
  27 def MockContainer(image_hash):
  28     return {'Id': MockDockerId(),
  29             'Image': image_hash['Id']}
  30
  31
  32 def MockImage(*, size=0, vsize=None, tags=[]):
  33     if vsize is None:
  34         vsize = random.randint(100, 2000000)
  35     return {'Id': MockDockerId(),
  36             'ParentId': MockDockerId(),
  37             'RepoTags': list(tags),
  38             'Size': size,
  39             'VirtualSize': vsize}
  40
  41
  42 class MockEvent(dict):
  43     ENCODING = 'utf-8'
  44     event_seq = itertools.count(1)
  45
  46     def __init__(self, status, docker_id=None, **event_data):
  47         if docker_id is None:
  48             docker_id = MockDockerId()
  49         super().__init__(self, **event_data)
  50         self['status'] = status
  51         self['id'] = docker_id
  52         self.setdefault('time', next(self.event_seq))
  53
  54     def encoded(self):
  55         return json.dumps(self).encode(self.ENCODING)
  56
  57
  58 class MockException(docker.errors.APIError):
  59
  60     def __init__(self, status_code):
  61         response = mock.Mock(name='response')
  62         response.status_code = status_code
  63         super().__init__("mock exception", response)
  64
  65
  66 class DockerImageTestCase(unittest.TestCase):
  67
  68     def test_used_at_sets_last_used(self):
  69         image = cleaner.DockerImage(MockImage())
  70         image.used_at(5)
  71         self.assertEqual(5, image.last_used)
  72
  73     def test_used_at_moves_forward(self):
  74         image = cleaner.DockerImage(MockImage())
  75         image.used_at(6)
  76         image.used_at(8)
  77         self.assertEqual(8, image.last_used)
  78
  79     def test_used_at_does_not_go_backward(self):
  80         image = cleaner.DockerImage(MockImage())
  81         image.used_at(4)
  82         image.used_at(2)
  83         self.assertEqual(4, image.last_used)
  84
  85
  86 class DockerImagesTestCase(unittest.TestCase):
  87
  88     def setUp(self):
  89         self.mock_images = []
  90
  91     def setup_mock_images(self, *vsizes):
  92         self.mock_images.extend(MockImage(vsize=vsize) for vsize in vsizes)
  93
  94     def setup_images(self, *vsizes, target_size=1000000):
  95         self.setup_mock_images(*vsizes)
  96         images = cleaner.DockerImages(target_size)
  97         for image in self.mock_images:
  98             images.add_image(image)
  99         return images
 100
 101     def test_has_image(self):
 102         images = self.setup_images(None)
 103         self.assertTrue(images.has_image(self.mock_images[0]['Id']))
 104         self.assertFalse(images.has_image(MockDockerId()))
 105
 106     def test_del_image(self):
 107         images = self.setup_images(None)
 108         images.del_image(self.mock_images[0]['Id'])
 109         self.assertFalse(images.has_image(self.mock_images[0]['Id']))
 110
 111     def test_del_nonexistent_image(self):
 112         images = self.setup_images(None)
 113         images.del_image(MockDockerId())
 114         self.assertTrue(images.has_image(self.mock_images[0]['Id']))
 115
 116     def test_one_image_always_kept(self):
 117         # When crunch-job starts a job, it makes sure each compute node
 118         # has the Docker image loaded, then it runs all the tasks with
 119         # the assumption the image is on each node.  As long as that's
 120         # true, the cleaner should avoid removing every installed image:
 121         # crunch-job might be counting on the most recent one to be
 122         # available, even if it's not currently in use.
 123         images = self.setup_images(None, None, target_size=1)
 124         for use_time, image in enumerate(self.mock_images, 1):
 125             user = MockContainer(image)
 126             images.add_user(user, use_time)
 127             images.end_user(user['Id'])
 128         self.assertEqual([self.mock_images[0]['Id']],
 129                          list(images.should_delete()))
 130
 131     def test_images_under_target_not_deletable(self):
 132         # The images are used in this order.  target_size is set so it
 133         # could hold the largest image, but not after the most recently
 134         # used image is kept; then we have to fall back to the previous one.
 135         images = self.setup_images(20, 30, 40, 10, target_size=45)
 136         for use_time, image in enumerate(self.mock_images, 1):
 137             user = MockContainer(image)
 138             images.add_user(user, use_time)
 139             images.end_user(user['Id'])
 140         self.assertEqual([self.mock_images[ii]['Id'] for ii in [0, 2]],
 141                          list(images.should_delete()))
 142
 143     def test_images_in_use_not_deletable(self):
 144         images = self.setup_images(None, None, target_size=1)
 145         users = [MockContainer(image) for image in self.mock_images]
 146         images.add_user(users[0], 1)
 147         images.add_user(users[1], 2)
 148         images.end_user(users[1]['Id'])
 149         self.assertEqual([self.mock_images[1]['Id']],
 150                          list(images.should_delete()))
 151
 152     def test_image_deletable_after_unused(self):
 153         images = self.setup_images(None, None, target_size=1)
 154         users = [MockContainer(image) for image in self.mock_images]
 155         images.add_user(users[0], 1)
 156         images.add_user(users[1], 2)
 157         images.end_user(users[0]['Id'])
 158         self.assertEqual([self.mock_images[0]['Id']],
 159                          list(images.should_delete()))
 160
 161     def test_image_not_deletable_if_user_restarts(self):
 162         images = self.setup_images(None, target_size=1)
 163         user = MockContainer(self.mock_images[-1])
 164         images.add_user(user, 1)
 165         images.end_user(user['Id'])
 166         images.add_user(user, 2)
 167         self.assertEqual([], list(images.should_delete()))
 168
 169     def test_image_not_deletable_if_any_user_remains(self):
 170         images = self.setup_images(None, target_size=1)
 171         users = [MockContainer(self.mock_images[0]) for ii in range(2)]
 172         images.add_user(users[0], 1)
 173         images.add_user(users[1], 2)
 174         images.end_user(users[0]['Id'])
 175         self.assertEqual([], list(images.should_delete()))
 176
 177     def test_image_deletable_after_all_users_end(self):
 178         images = self.setup_images(None, None, target_size=1)
 179         users = [MockContainer(self.mock_images[ii]) for ii in [0, 1, 1]]
 180         images.add_user(users[0], 1)
 181         images.add_user(users[1], 2)
 182         images.add_user(users[2], 3)
 183         images.end_user(users[1]['Id'])
 184         images.end_user(users[2]['Id'])
 185         self.assertEqual([self.mock_images[-1]['Id']],
 186                          list(images.should_delete()))
 187
 188     def test_images_suggested_for_deletion_by_lru(self):
 189         images = self.setup_images(10, 10, 10, target_size=1)
 190         users = [MockContainer(image) for image in self.mock_images]
 191         images.add_user(users[0], 3)
 192         images.add_user(users[1], 1)
 193         images.add_user(users[2], 2)
 194         for user in users:
 195             images.end_user(user['Id'])
 196         self.assertEqual([self.mock_images[ii]['Id'] for ii in [1, 2]],
 197                          list(images.should_delete()))
 198
 199     def test_adding_user_without_image_does_not_implicitly_add_image(self):
 200         images = self.setup_images(10)
 201         images.add_user(MockContainer(MockImage()), 1)
 202         self.assertEqual([], list(images.should_delete()))
 203
 204     def test_nonexistent_user_removed(self):
 205         images = self.setup_images()
 206         images.end_user('nonexistent')
 207         # No exception should be raised.
 208
 209     def test_del_image_effective_with_users_present(self):
 210         images = self.setup_images(None, target_size=1)
 211         user = MockContainer(self.mock_images[0])
 212         images.add_user(user, 1)
 213         images.del_image(self.mock_images[0]['Id'])
 214         images.end_user(user['Id'])
 215         self.assertEqual([], list(images.should_delete()))
 216
 217     def setup_from_daemon(self, *vsizes, target_size=1500000):
 218         self.setup_mock_images(*vsizes)
 219         docker_client = mock.MagicMock(name='docker_client')
 220         docker_client.images.return_value = iter(self.mock_images)
 221         return cleaner.DockerImages.from_daemon(target_size, docker_client)
 222
 223     def test_images_loaded_from_daemon(self):
 224         images = self.setup_from_daemon(None, None)
 225         for image in self.mock_images:
 226             self.assertTrue(images.has_image(image['Id']))
 227
 228     def test_target_size_set_from_daemon(self):
 229         images = self.setup_from_daemon(20, 10, 5, target_size=15)
 230         user = MockContainer(self.mock_images[-1])
 231         images.add_user(user, 1)
 232         self.assertEqual([self.mock_images[0]['Id']],
 233                          list(images.should_delete()))
 234
 235
 236 class DockerImageUseRecorderTestCase(unittest.TestCase):
 237     TEST_CLASS = cleaner.DockerImageUseRecorder
 238     TEST_CLASS_INIT_KWARGS = {}
 239
 240     def setUp(self):
 241         self.images = mock.MagicMock(name='images')
 242         self.docker_client = mock.MagicMock(name='docker_client')
 243         self.events = []
 244         self.recorder = self.TEST_CLASS(self.images, self.docker_client,
 245                                         self.encoded_events, **self.TEST_CLASS_INIT_KWARGS)
 246
 247     @property
 248     def encoded_events(self):
 249         return (event.encoded() for event in self.events)
 250
 251     def test_unknown_events_ignored(self):
 252         self.events.append(MockEvent('mock!event'))
 253         self.recorder.run()
 254         # No exception should be raised.
 255
 256     def test_fetches_container_on_create(self):
 257         self.events.append(MockEvent('create'))
 258         self.recorder.run()
 259         self.docker_client.inspect_container.assert_called_with(
 260             self.events[0]['id'])
 261
 262     def test_adds_user_on_container_create(self):
 263         self.events.append(MockEvent('create'))
 264         self.recorder.run()
 265         self.images.add_user.assert_called_with(
 266             self.docker_client.inspect_container(), self.events[0]['time'])
 267
 268     def test_unknown_image_handling(self):
 269         # The use recorder should not fetch any images.
 270         self.events.append(MockEvent('create'))
 271         self.recorder.run()
 272         self.assertFalse(self.docker_client.inspect_image.called)
 273
 274     def test_unfetchable_containers_ignored(self):
 275         self.events.append(MockEvent('create'))
 276         self.docker_client.inspect_container.side_effect = MockException(404)
 277         self.recorder.run()
 278         self.assertFalse(self.images.add_user.called)
 279
 280     def test_ends_user_on_container_destroy(self):
 281         self.events.append(MockEvent('destroy'))
 282         self.recorder.run()
 283         self.images.end_user.assert_called_with(self.events[0]['id'])
 284
 285
 286 class DockerImageCleanerTestCase(DockerImageUseRecorderTestCase):
 287     TEST_CLASS = cleaner.DockerImageCleaner
 288
 289     def test_unknown_image_handling(self):
 290         # The image cleaner should fetch and record new images.
 291         self.images.has_image.return_value = False
 292         self.events.append(MockEvent('create'))
 293         self.recorder.run()
 294         self.docker_client.inspect_image.assert_called_with(
 295             self.docker_client.inspect_container()['Image'])
 296         self.images.add_image.assert_called_with(
 297             self.docker_client.inspect_image())
 298
 299     def test_unfetchable_images_ignored(self):
 300         self.images.has_image.return_value = False
 301         self.docker_client.inspect_image.side_effect = MockException(404)
 302         self.events.append(MockEvent('create'))
 303         self.recorder.run()
 304         self.docker_client.inspect_image.assert_called_with(
 305             self.docker_client.inspect_container()['Image'])
 306         self.assertFalse(self.images.add_image.called)
 307
 308     def test_deletions_after_destroy(self):
 309         delete_id = MockDockerId()
 310         self.images.should_delete.return_value = [delete_id]
 311         self.events.append(MockEvent('destroy'))
 312         self.recorder.run()
 313         self.docker_client.remove_image.assert_called_with(delete_id)
 314         self.images.del_image.assert_called_with(delete_id)
 315
 316     def test_failed_deletion_handling(self):
 317         delete_id = MockDockerId()
 318         self.images.should_delete.return_value = [delete_id]
 319         self.docker_client.remove_image.side_effect = MockException(500)
 320         self.events.append(MockEvent('destroy'))
 321         self.recorder.run()
 322         self.docker_client.remove_image.assert_called_with(delete_id)
 323         self.assertFalse(self.images.del_image.called)
 324
 325
 326 class DockerContainerCleanerTestCase(DockerImageUseRecorderTestCase):
 327     TEST_CLASS = cleaner.DockerImageCleaner
 328     TEST_CLASS_INIT_KWARGS = {'remove_containers_onexit': True}
 329
 330     def test_container_deletion_deletes_volumes(self):
 331         cid = MockDockerId()
 332         self.events.append(MockEvent('die', docker_id=cid))
 333         self.recorder.run()
 334         self.docker_client.remove_container.assert_called_with(cid, v=True)
 335
 336     @mock.patch('arvados_docker.cleaner.logger')
 337     def test_failed_container_deletion_handling(self, mockLogger):
 338         cid = MockDockerId()
 339         self.docker_client.remove_container.side_effect = MockException(500)
 340         self.events.append(MockEvent('die', docker_id=cid))
 341         self.recorder.run()
 342         self.docker_client.remove_container.assert_called_with(cid, v=True)
 343         self.assertEqual("Failed to remove container %s: %s",
 344                          mockLogger.warning.call_args[0][0])
 345         self.assertEqual(cid,
 346                          mockLogger.warning.call_args[0][1])
 347
 348
 349 class HumanSizeTestCase(unittest.TestCase):
 350
 351     def check(self, human_str, count, exp):
 352         self.assertEqual(count * (1024 ** exp),
 353                          cleaner.human_size(human_str))
 354
 355     def test_bytes(self):
 356         self.check('1', 1, 0)
 357         self.check('82', 82, 0)
 358
 359     def test_kibibytes(self):
 360         self.check('2K', 2, 1)
 361         self.check('3k', 3, 1)
 362
 363     def test_mebibytes(self):
 364         self.check('4M', 4, 2)
 365         self.check('5m', 5, 2)
 366
 367     def test_gibibytes(self):
 368         self.check('6G', 6, 3)
 369         self.check('7g', 7, 3)
 370
 371     def test_tebibytes(self):
 372         self.check('8T', 8, 4)
 373         self.check('9t', 9, 4)
 374
 375
 376 class RunTestCase(unittest.TestCase):
 377
 378     def setUp(self):
 379         self.config = cleaner.default_config()
 380         self.config['Quota'] = 1000000
 381         self.docker_client = mock.MagicMock(name='docker_client')
 382
 383     def test_run(self):
 384         test_start_time = int(time.time())
 385         self.docker_client.events.return_value = []
 386         cleaner.run(self.config, self.docker_client)
 387         self.assertEqual(2, self.docker_client.events.call_count)
 388         event_kwargs = [args[1] for args in
 389                         self.docker_client.events.call_args_list]
 390         self.assertIn('since', event_kwargs[0])
 391         self.assertIn('until', event_kwargs[0])
 392         self.assertLessEqual(test_start_time, event_kwargs[0]['until'])
 393         self.assertIn('since', event_kwargs[1])
 394         self.assertEqual(event_kwargs[0]['until'], event_kwargs[1]['since'])
 395
 396
 397 @mock.patch('docker.APIClient', name='docker_client')
 398 @mock.patch('arvados_docker.cleaner.run', name='cleaner_run')
 399 class MainTestCase(unittest.TestCase):
 400
 401     def test_client_api_version(self, run_mock, docker_client):
 402         with tempfile.NamedTemporaryFile(mode='wt') as cf:
 403             cf.write('{"Quota":"1000T"}')
 404             cf.flush()
 405             cleaner.main(['--config', cf.name])
 406         self.assertEqual(1, docker_client.call_count)
 407         # We are standardized on Docker API version 1.35.
 408         # See DockerAPIVersion in lib/crunchrun/docker.go.
 409         self.assertEqual('1.35',
 410                          docker_client.call_args[1].get('version'))
 411         self.assertEqual(1, run_mock.call_count)
 412         self.assertIs(run_mock.call_args[0][1], docker_client())
 413
 414
 415 class ConfigTestCase(unittest.TestCase):
 416
 417     def test_load_config(self):
 418         with tempfile.NamedTemporaryFile(mode='wt') as cf:
 419             cf.write(
 420                 '{"Quota":"1000T", "RemoveStoppedContainers":"always", "Verbose":2}')
 421             cf.flush()
 422             config = cleaner.load_config(['--config', cf.name])
 423         self.assertEqual(1000 << 40, config['Quota'])
 424         self.assertEqual("always", config['RemoveStoppedContainers'])
 425         self.assertEqual(2, config['Verbose'])
 426
 427     def test_args_override_config(self):
 428         with tempfile.NamedTemporaryFile(mode='wt') as cf:
 429             cf.write(
 430                 '{"Quota":"1000T", "RemoveStoppedContainers":"always", "Verbose":2}')
 431             cf.flush()
 432             config = cleaner.load_config([
 433                 '--config', cf.name,
 434                 '--quota', '1G',
 435                 '--remove-stopped-containers', 'never',
 436                 '--verbose',
 437             ])
 438         self.assertEqual(1 << 30, config['Quota'])
 439         self.assertEqual('never', config['RemoveStoppedContainers'])
 440         self.assertEqual(1, config['Verbose'])
 441
 442     def test_args_no_config(self):
 443         self.assertEqual(False, os.path.exists(cleaner.DEFAULT_CONFIG_FILE))
 444         config = cleaner.load_config(['--quota', '1G'])
 445         self.assertEqual(1 << 30, config['Quota'])
 446
 447
 448 class ContainerRemovalTestCase(unittest.TestCase):
 449     LIFECYCLE = ['create', 'attach', 'start', 'resize', 'die', 'destroy']
 450
 451     def setUp(self):
 452         self.config = cleaner.default_config()
 453         self.docker_client = mock.MagicMock(name='docker_client')
 454         self.existingCID = MockDockerId()
 455         self.docker_client.containers.return_value = [{
 456             'Id': self.existingCID,
 457             'Status': 'Exited (0) 6 weeks ago',
 458         }, {
 459             # If docker_client.containers() returns non-exited
 460             # containers for some reason, do not remove them.
 461             'Id': MockDockerId(),
 462             'Status': 'Running',
 463         }]
 464         self.newCID = MockDockerId()
 465         self.docker_client.events.return_value = [
 466             MockEvent(e, docker_id=self.newCID).encoded()
 467             for e in self.LIFECYCLE]
 468
 469     def test_remove_onexit(self):
 470         self.config['RemoveStoppedContainers'] = 'onexit'
 471         cleaner.run(self.config, self.docker_client)
 472         self.docker_client.remove_container.assert_called_once_with(
 473             self.newCID, v=True)
 474
 475     def test_remove_always(self):
 476         self.config['RemoveStoppedContainers'] = 'always'
 477         cleaner.run(self.config, self.docker_client)
 478         self.docker_client.remove_container.assert_any_call(
 479             self.existingCID, v=True)
 480         self.docker_client.remove_container.assert_any_call(
 481             self.newCID, v=True)
 482         self.assertEqual(2, self.docker_client.remove_container.call_count)
 483
 484     def test_remove_never(self):
 485         self.config['RemoveStoppedContainers'] = 'never'
 486         cleaner.run(self.config, self.docker_client)
 487         self.assertEqual(0, self.docker_client.remove_container.call_count)
 488
 489     def test_container_exited_between_subscribe_events_and_check_existing(self):
 490         self.config['RemoveStoppedContainers'] = 'always'
 491         self.docker_client.events.return_value = [
 492             MockEvent(e, docker_id=self.existingCID).encoded()
 493             for e in ['die', 'destroy']]
 494         cleaner.run(self.config, self.docker_client)
 495         # Subscribed to events before getting the list of existing
 496         # exited containers?
 497         self.docker_client.assert_has_calls([
 498             mock.call.events(since=mock.ANY),
 499             mock.call.containers(filters={'status': 'exited'})])
 500         # Asked to delete the container twice?
 501         self.docker_client.remove_container.assert_has_calls(
 502             [mock.call(self.existingCID, v=True)] * 2)
 503         self.assertEqual(2, self.docker_client.remove_container.call_count)