allow modifying the schedule at runtime at recalculate heap when needed

Jeroeny · Jeroeny · commit 92fa7109962a · 2023-09-04T14:14:48.000+02:00
diff --git a/src/Symfony/Component/Scheduler/Generator/MessageGenerator.php b/src/Symfony/Component/Scheduler/Generator/MessageGenerator.php
@@ -19,7 +19,7 @@
 
 final class MessageGenerator implements MessageGeneratorInterface
 {
-    private Schedule $schedule;
+    private ?Schedule $schedule = null;
     private TriggerHeap $triggerHeap;
     private ?\DateTimeImmutable $waitUntil;
 
@@ -36,6 +36,12 @@ public function getMessages(): \Generator
     {
         $checkpoint = $this->checkpoint();
 
+        if ($this->schedule?->restart()) {
+            unset($this->triggerHeap);
+            $this->waitUntil = new \DateTimeImmutable('@0');
+            $this->schedule->restart(false);
+        }
+
         if (!$this->waitUntil
             || $this->waitUntil > ($now = $this->clock->now())
             || !$checkpoint->acquire($now)
diff --git a/src/Symfony/Component/Scheduler/Schedule.php b/src/Symfony/Component/Scheduler/Schedule.php
@@ -22,19 +22,49 @@ final class Schedule implements ScheduleProviderInterface
     private ?LockInterface $lock = null;
     private ?CacheInterface $state = null;
 
+    private bool $restart = false;
+
+    public static function with(RecurringMessage $message, RecurringMessage ...$messages): static
+    {
+        return static::doAdd(new self(), $message, ...$messages);
+    }
+
     /**
      * @return $this
      */
     public function add(RecurringMessage $message, RecurringMessage ...$messages): static
+    {
+        $this->restart(true);
+
+        return static::doAdd($this, $message, ...$messages);
+    }
+
+    private static function doAdd(self $schedule, RecurringMessage $message, RecurringMessage ...$messages): static
     {
         foreach ([$message, ...$messages] as $m) {
-            if (isset($this->messages[$m->getId()])) {
+            if (isset($schedule->messages[$m->getId()])) {
                 throw new LogicException('Duplicated schedule message.');
             }
 
-            $this->messages[$m->getId()] = $m;
+            $schedule->messages[$m->getId()] = $m;
         }
 
+        return $schedule;
+    }
+
+    public function remove(RecurringMessage $message): static
+    {
+        unset($this->messages[$message->getId()]);
+        $this->restart(true);
+
+        return $this;
+    }
+
+    public function clear(): static
+    {
+        $this->messages = [];
+        $this->restart(true);
+
         return $this;
     }
 
@@ -83,4 +113,9 @@ public function getSchedule(): static
     {
         return $this;
     }
+
+    public function restart(bool $restart = null): bool
+    {
+        return $this->restart = $restart ?? $this->restart;
+    }
 }
diff --git a/src/Symfony/Component/Scheduler/Tests/Generator/MessageGeneratorTest.php b/src/Symfony/Component/Scheduler/Tests/Generator/MessageGeneratorTest.php
@@ -95,6 +95,72 @@ public function getSchedule(): Schedule
         }
     }
 
+    public function testGetMessagesFromScheduleProviderWithRestart()
+    {
+        $first = (object) ['id' => 'first'];
+        $startTime = '22:12:00';
+        $runs = [
+            '22:12:00' => [],
+            '22:12:01' => [],
+            '22:13:00' => [$first],
+            '22:13:01' => [],
+        ];
+        $schedule = [[$first, '22:13:00', '22:14:00']];
+
+        // for referencing
+        $now = self::makeDateTime($startTime);
+
+        $clock = $this->createMock(ClockInterface::class);
+        $clock->method('now')->willReturnReference($now);
+
+        foreach ($schedule as $i => $s) {
+            if (\is_array($s)) {
+                $schedule[$i] = $this->createMessage(...$s);
+            }
+        }
+
+        $scheduleProvider = new class($schedule) implements ScheduleProviderInterface {
+            private Schedule $schedule;
+
+            public function __construct(array $schedule)
+            {
+                $this->schedule = Schedule::with(...$schedule);
+                $this->schedule->stateful(new ArrayAdapter());
+            }
+
+            public function getSchedule(): Schedule
+            {
+                return $this->schedule;
+            }
+
+            public function add(RecurringMessage $message): self
+            {
+                $this->schedule->add($message);
+
+                return $this;
+            }
+        };
+
+        $scheduler = new MessageGenerator($scheduleProvider, 'dummy', $clock);
+
+        // Warmup. The first run always returns nothing.
+        $this->assertSame([], iterator_to_array($scheduler->getMessages(), false));
+
+        $toAdd = (object) ['id' => 'added-after-start'];
+
+        foreach ($runs as $time => $expected) {
+            $now = self::makeDateTime($time);
+            $this->assertSame($expected, iterator_to_array($scheduler->getMessages(), false));
+        }
+
+        $scheduleProvider->add($this->createMessage($toAdd, '22:13:10', '22:13:11'));
+
+        $this->assertSame([], iterator_to_array($scheduler->getMessages(), false));
+
+        $now = self::makeDateTime('22:13:10');
+        $this->assertSame([$toAdd], iterator_to_array($scheduler->getMessages(), false));
+    }
+
     public function testYieldedContext()
     {
         // for referencing
