Passed
Pull Request — master (#135)
by Aldo
03:45
created

MaintenanceScheduler.shutdown()   A

Complexity

Conditions 2

Size

Total Lines 12
Code Lines 6

Duplication

Lines 0
Ratio 0 %

Code Coverage

Tests 6
CRAP Score 2

Importance

Changes 0
Metric Value
cc 2
eloc 6
nop 1
dl 0
loc 12
ccs 6
cts 6
cp 1
crap 2
rs 10
c 0
b 0
f 0
1
"""Module for handling the scheduled execution of maintenance windows."""
2 1
import pytz
3 1
from dataclasses import dataclass
4
5 1
from apscheduler.jobstores.base import JobLookupError
6 1
from apscheduler.schedulers.background import BackgroundScheduler
7 1
from apscheduler.schedulers.base import BaseScheduler
8
9
10 1
from .deployer import MaintenanceDeployer
11 1
from ..controllers import MaintenanceController
12 1
from ..models import (
13
    MaintenanceID,
14
    MaintenanceWindow,
15
    MaintenanceWindows,
16
    OverlapError,
17
    Status,
18
)
19
20 1
from kytos.core import log
21
22 1
@dataclass
23 1
class MaintenanceStart:
24
    """
25
    Callable used for starting maintenance windows
26
    """
27 1
    maintenance_scheduler: 'MaintenanceScheduler'
28 1
    mw_id: MaintenanceID
29
30 1
    def __call__(self):
31 1
        self.maintenance_scheduler.start_maintenance(self.mw_id)
32
33
34 1
@dataclass
35 1
class MaintenanceEnd:
36
    """
37
    Callable used for ending maintenance windows
38
    """
39 1
    maintenance_scheduler: 'MaintenanceScheduler'
40 1
    mw_id: MaintenanceID
41
42 1
    def __call__(self):
43 1
        self.maintenance_scheduler.end_maintenance(self.mw_id)
44
45 1
@dataclass
46 1
class MaintenanceScheduler:
47
    """Class for scheduling maintenance windows."""
48 1
    deployer: MaintenanceDeployer
49 1
    db_controller: MaintenanceController
50 1
    scheduler: BaseScheduler
51
52 1
    @classmethod
53 1
    def new_scheduler(cls, deployer: MaintenanceDeployer):
54
        """
55
        Creates a new scheduler from the given MaintenanceDeployer
56
        """
57
        scheduler = BackgroundScheduler(timezone=pytz.utc)
58
        db_controller = MaintenanceController()
59
        db_controller.bootstrap_indexes()
60
        instance = cls(deployer, db_controller, scheduler)
61
        return instance
62
63 1
    def start(self):
64
        """
65
        Begin running the scheduler.
66
        """
67 1
        self.db_controller.prepare_start()
68
69
        # Populate the scheduler with all pending tasks
70 1
        windows = self.db_controller.get_unfinished_windows()
71 1
        for window in windows:
72 1
            if window.status == Status.RUNNING:
73 1
                self.deployer.start_mw(window)
74 1
            self._schedule(window)
75
76
        # Start the scheduler
77 1
        self.scheduler.start()
78
79 1
    def shutdown(self):
80
        """
81
        Stop running the scheduler.
82
        """
83 1
        windows = self.db_controller.get_windows()
84
85
        # Depopulate the scheduler
86 1
        for window in windows:
87 1
            self._unschedule(window)
88
89 1
        self.scheduler.remove_all_jobs()
90 1
        self.scheduler.shutdown()
91
92 1
    def start_maintenance(self, mw_id: MaintenanceID):
93
        """Begins executing the maintenance window
94
        """
95
        # Get Maintenance from DB and Update
96 1
        window = self.db_controller.start_window(mw_id)
97
98
        # Activate Running
99 1
        self.deployer.start_mw(window)
100
101
        # Schedule next task
102 1
        self._schedule(window)
103
104 1
    def end_maintenance(self, mw_id: MaintenanceID):
105
        """Ends execution of the maintenance window
106
        """
107
        # Get Maintenance from DB
108 1
        window = self.db_controller.end_window(mw_id)
109
110
        # Set to Ending
111 1
        self.deployer.end_mw(window)
112
113 1
    def end_maintenance_early(self, mw_id: MaintenanceID):
114
        """Ends execution of the maintenance window early
115
        """
116
        # Get Maintenance from DB
117
        window = self.db_controller.end_window(mw_id)
118
119
        # Unschedule tasks
120
        self._unschedule(window)
121
122 1
    def add(self, window: MaintenanceWindow, force=False):
123
        """Add jobs to start and end a maintenance window."""
124
        overlapping_windows = self.db_controller.check_overlap(window, force)
125
        if overlapping_windows:
126
            raise OverlapError(window, overlapping_windows)
127
128
        # Add window to DB
129
        self.db_controller.insert_window(window)
130
131
        # Schedule next task
132
        self._schedule(window)
133
134 1
    def update(self, window: MaintenanceWindow):
135
        """Update an existing Maintenance Window."""
136
137
        # Update window
138 1
        self.db_controller.update_window(window)
139
140
        # Reschedule any pending tasks
141 1
        self._reschedule(window)
142
143 1
    def remove(self, mw_id: MaintenanceID):
144
        """Remove jobs that start and end a maintenance window."""
145
        # Get Maintenance from DB
146
        window = self.db_controller.get_window(mw_id)
147
148
        # Remove from schedule
149
        self._unschedule(window)
150
151
        # Remove from DB
152
        self.db_controller.remove_window(mw_id)
153
154 1
    def _schedule(self, window: MaintenanceWindow):
155 1
        log.info(f'Scheduling "{window.id}"')
156 1
        if window.status == Status.PENDING:
157 1
            self.scheduler.add_job(
158
                MaintenanceStart(self, window.id),
159
                'date',
160
                id=f'{window.id}-start',
161
                run_date=window.start
162
            )
163 1
            log.info(f'Scheduled "{window.id}" start at {window.start}')
164 1
        if window.status == Status.RUNNING:
165 1
            self.scheduler.add_job(
166
                MaintenanceEnd(self, window.id),
167
                'date',
168
                id=f'{window.id}-end',
169
                run_date=window.end
170
            )
171 1
            log.info(f'Scheduled "{window.id}" end at {window.end}')
172
173 1
    def _reschedule(self, window: MaintenanceWindow):
174 1
        log.info(f'Rescheduling "{window.id}"')
175 1
        try:
176 1
            self.scheduler.remove_job(
177
                f'{window.id}-start',
178
            )
179 1
            self.scheduler.add_job(
180
                MaintenanceStart(self, window.id),
181
                'date',
182
                id=f'{window.id}-start',
183
                run_date=window.start
184
            )
185 1
            log.info(f'Rescheduled "{window.id}" start to {window.start}')
186
        except JobLookupError:
187
            log.info(f'Could not reschedule "{window.id}" start, no start job')
188 1
        try:
189 1
            self.scheduler.remove_job(
190
                f'{window.id}-end',
191
            )
192 1
            self.scheduler.add_job(
193
                MaintenanceEnd(self, window.id),
194
                'date',
195
                id=f'{window.id}-end',
196
                run_date=window.end
197
            )
198 1
            log.info(f'Rescheduled "{window.id}" end to {window.end}')
199
        except JobLookupError:
200
            log.info(f'Could not reschedule "{window.id}" end, no end job')
201
202 1
    def _unschedule(self, window: MaintenanceWindow):
203
        """Remove maintenance events from scheduler.
204
        Does not update DB, due to being
205
        primarily for shutdown startup cases.
206
        """
207 1
        started = False
208 1
        ended = False
209 1
        try:
210 1
            self.scheduler.remove_job(f'{window.id}-start')
211 1
        except JobLookupError:
212 1
            started = True
213 1
            log.info(f'Job to start "{window.id}" already removed.')
214 1
        try:
215 1
            self.scheduler.remove_job(f'{window.id}-end')
216 1
        except JobLookupError:
217 1
            ended = True
218 1
            log.info(f'Job to end "{window.id}" already removed.')
219 1
        if started and not ended:
220 1
            self.deployer.end_mw(window)
221
222 1
    def get_maintenance(self, mw_id: MaintenanceID) -> MaintenanceWindow:
223
        """Get a single maintenance by id"""
224
        return self.db_controller.get_window(mw_id)
225
226 1
    def list_maintenances(self) -> MaintenanceWindows:
227
        """Returns a list of all maintenances"""
228
        return self.db_controller.get_windows()
229