summaryrefslogtreecommitdiffstats
path: root/stator/models.py
blob: 5257ac9fdee5bec188dccdb89819184bb923d839 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
import datetime
import pprint
import traceback
from typing import ClassVar, List, Optional, Type, Union, cast

from asgiref.sync import sync_to_async
from django.db import models, transaction
from django.utils import timezone
from django.utils.functional import classproperty

from core import exceptions
from stator.graph import State, StateGraph


class StateField(models.CharField):
    """
    A special field that automatically gets choices from a state graph
    """

    def __init__(self, graph: Type[StateGraph], **kwargs):
        # Sensible default for state length
        kwargs.setdefault("max_length", 100)
        # Add choices and initial
        self.graph = graph
        kwargs["choices"] = self.graph.choices
        kwargs["default"] = self.graph.initial_state.name
        super().__init__(**kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super().deconstruct()
        kwargs["graph"] = self.graph
        return name, path, args, kwargs

    def get_prep_value(self, value):
        if isinstance(value, State):
            return value.name
        return value


class StatorModel(models.Model):
    """
    A model base class that has a state machine backing it, with tasks to work
    out when to move the state to the next one.

    You need to provide a "state" field as an instance of StateField on the
    concrete model yourself.
    """

    # If this row is up for transition attempts (which it always is on creation!)
    state_ready = models.BooleanField(default=True)

    # When the state last actually changed, or the date of instance creation
    state_changed = models.DateTimeField(auto_now_add=True)

    # When the last state change for the current state was attempted
    # (and not successful, as this is cleared on transition)
    state_attempted = models.DateTimeField(blank=True, null=True)

    # If a lock is out on this row, when it is locked until
    # (we don't identify the lock owner, as there's no heartbeats)
    state_locked_until = models.DateTimeField(null=True, blank=True)

    # Collection of subclasses of us
    subclasses: ClassVar[List[Type["StatorModel"]]] = []

    class Meta:
        abstract = True

    def __init_subclass__(cls) -> None:
        if cls is not StatorModel:
            cls.subclasses.append(cls)

    @classproperty
    def state_graph(cls) -> Type[StateGraph]:
        return cls._meta.get_field("state").graph

    @property
    def state_age(self) -> int:
        return (timezone.now() - self.state_changed).total_seconds()

    @classmethod
    async def atransition_schedule_due(cls, now=None) -> models.QuerySet:
        """
        Finds instances of this model that need to run and schedule them.
        """
        q = models.Q()
        for state in cls.state_graph.states.values():
            state = cast(State, state)
            if not state.externally_progressed:
                q = q | models.Q(
                    (
                        models.Q(
                            state_attempted__lte=timezone.now()
                            - datetime.timedelta(
                                seconds=cast(float, state.try_interval)
                            )
                        )
                        | models.Q(state_attempted__isnull=True)
                    ),
                    state=state.name,
                )
        await cls.objects.filter(q).aupdate(state_ready=True)

    @classmethod
    def transition_get_with_lock(
        cls, number: int, lock_expiry: datetime.datetime
    ) -> List["StatorModel"]:
        """
        Returns up to `number` tasks for execution, having locked them.
        """
        with transaction.atomic():
            selected = list(
                cls.objects.filter(
                    state_locked_until__isnull=True,
                    state_ready=True,
                    state__in=cls.state_graph.automatic_states,
                )[:number].select_for_update()
            )
            cls.objects.filter(pk__in=[i.pk for i in selected]).update(
                state_locked_until=lock_expiry
            )
        return selected

    @classmethod
    async def atransition_get_with_lock(
        cls, number: int, lock_expiry: datetime.datetime
    ) -> List["StatorModel"]:
        return await sync_to_async(cls.transition_get_with_lock)(number, lock_expiry)

    @classmethod
    async def atransition_clean_locks(cls):
        await cls.objects.filter(state_locked_until__lte=timezone.now()).aupdate(
            state_locked_until=None
        )

    def transition_schedule(self):
        """
        Adds this instance to the queue to get its state transition attempted.

        The scheduler will call this, but you can also call it directly if you
        know it'll be ready and want to lower latency.
        """
        self.state_ready = True
        self.save()

    async def atransition_attempt(self) -> Optional[State]:
        """
        Attempts to transition the current state by running its handler(s).
        """
        current_state = self.state_graph.states[self.state]
        # If it's a manual progression state don't even try
        # We shouldn't really be here in this case, but it could be a race condition
        if current_state.externally_progressed:
            print(
                f"Warning: trying to progress externally progressed state {self.state}!"
            )
            return None
        try:
            next_state = await current_state.handler(self)
        except BaseException as e:
            await StatorError.acreate_from_instance(self, e)
            await exceptions.acapture_exception(e)
            traceback.print_exc()
        else:
            if next_state:
                # Ensure it's a State object
                if isinstance(next_state, str):
                    next_state = self.state_graph.states[next_state]
                # Ensure it's a child
                if next_state not in current_state.children:
                    raise ValueError(
                        f"Cannot transition from {current_state} to {next_state} - not a declared transition"
                    )
                await self.atransition_perform(next_state)
                return next_state
        await self.__class__.objects.filter(pk=self.pk).aupdate(
            state_attempted=timezone.now(),
            state_locked_until=None,
            state_ready=False,
        )
        return None

    def transition_perform(self, state: Union[State, str]):
        """
        Transitions the instance to the given state name, forcibly.
        """
        if isinstance(state, State):
            state = state.name
        if state not in self.state_graph.states:
            raise ValueError(f"Invalid state {state}")
        # See if it's ready immediately (if not, delay until first try_interval)
        if self.state_graph.states[state].attempt_immediately:
            self.__class__.objects.filter(pk=self.pk).update(
                state=state,
                state_changed=timezone.now(),
                state_attempted=None,
                state_locked_until=None,
                state_ready=True,
            )
        else:
            self.__class__.objects.filter(pk=self.pk).update(
                state=state,
                state_changed=timezone.now(),
                state_attempted=timezone.now(),
                state_locked_until=None,
                state_ready=False,
            )

    atransition_perform = sync_to_async(transition_perform)


class StatorError(models.Model):
    """
    Tracks any errors running the transitions.
    Meant to be cleaned out regularly. Should probably be a log.
    """

    # appname.modelname (lowercased) label for the model this represents
    model_label = models.CharField(max_length=200)

    # The primary key of that model (probably int or str)
    instance_pk = models.CharField(max_length=200)

    # The state we were on
    state = models.CharField(max_length=200)

    # When it happened
    date = models.DateTimeField(auto_now_add=True)

    # Error name
    error = models.TextField()

    # Error details
    error_details = models.TextField(blank=True, null=True)

    @classmethod
    async def acreate_from_instance(
        cls,
        instance: StatorModel,
        exception: Optional[BaseException] = None,
    ):
        detail = traceback.format_exc()
        if exception and len(exception.args) > 1:
            detail += "\n\n" + "\n\n".join(
                pprint.pformat(arg) for arg in exception.args
            )

        return await cls.objects.acreate(
            model_label=instance._meta.label_lower,
            instance_pk=str(instance.pk),
            state=instance.state,
            error=str(exception),
            error_details=detail,
        )