NixOS/auto-upgrade: offer the possibility to define a reboot window during which the system may be automatically rebooted

Some systems should not be rebooted at just any time. If the upgrade process takes too long, for instance because of a
slow internet connection, or if the upgrade service is ran during production hours, we want to allow to define a window
outside of which a reboot will not be performed.
The system will then reboot on the next run of the upgrade service which finishes inside the reboot window.

E.g. we can run the update service twice per week, once during the night and once during the day, but reboots are only
allowed during the night. By doing so, a system that is usually shut down during the night will still receive updates
and systems that are turned on 24/7 can be rebooted outside of production hours.

Co-authored-by: Silvan Mosberger <github@infinisil.com>

1 files changed, 71 insertions, 12 deletions

diff --git a/nixos/modules/tasks/auto-upgrade.nix b/nixos/modules/tasks/auto-upgrade.nix
index b931b27ad81..1404dcbaf7c 100644
--- a/nixos/modules/tasks/auto-upgrade.nix
+++ b/nixos/modules/tasks/auto-upgrade.nix
@@ -80,6 +80,7 @@ in {
           Reboot the system into the new generation instead of a switch
           if the new generation uses a different kernel, kernel modules
           or initrd than the booted system.
+          See <option>rebootWindow</option> for configuring the times at which a reboot is allowed.
         '';
       };
 
@@ -96,6 +97,32 @@ in {
         '';
       };
 
+      rebootWindow = mkOption {
+        description = ''
+          Define a lower and upper time value (in HH:MM format) which
+          constitute a time window during which reboots are allowed after an upgrade.
+          This option only has an effect when <option>allowReboot</option> is enabled.
+          The default value of <literal>null</literal> means that reboots are allowed at any time.
+        '';
+        default = null;
+        example = { lower = "01:00"; upper = "05:00"; };
+        type = with types; nullOr (submodule {
+          options = {
+            lower = mkOption {
+              description = "Lower limit of the reboot window";
+              type = types.strMatching "[[:digit:]]{2}:[[:digit:]]{2}";
+              example = "01:00";
+            };
+
+            upper = mkOption {
+              description = "Upper limit of the reboot window";
+              type = types.strMatching "[[:digit:]]{2}:[[:digit:]]{2}";
+              example = "05:00";
+            };
+          };
+        });
+      };
+
     };
 
   };
@@ -110,12 +137,10 @@ in {
     }];
 
     system.autoUpgrade.flags = (if cfg.flake == null then
-        [ "--no-build-output" ] ++ (if cfg.channel == null then
-          [ "--upgrade" ]
-        else [
+        [ "--no-build-output" ] ++ optionals (cfg.channel != null) [
           "-I"
           "nixpkgs=${cfg.channel}/nixexprs.tar.xz"
-        ])
+        ]
       else
         [ "--flake ${cfg.flake}" ]);
 
@@ -143,19 +168,52 @@ in {
       ];
 
       script = let
-        nixos-rebuild =
-          "${config.system.build.nixos-rebuild}/bin/nixos-rebuild";
+        nixos-rebuild = "${config.system.build.nixos-rebuild}/bin/nixos-rebuild";
+        date     = "${pkgs.coreutils}/bin/date";
+        readlink = "${pkgs.coreutils}/bin/readlink";
+        shutdown = "${pkgs.systemd}/bin/shutdown";
+        upgradeFlag = optional (cfg.channel == null) "--upgrade";
       in if cfg.allowReboot then ''
-        ${nixos-rebuild} boot ${toString cfg.flags}
-        booted="$(readlink /run/booted-system/{initrd,kernel,kernel-modules})"
-        built="$(readlink /nix/var/nix/profiles/system/{initrd,kernel,kernel-modules})"
-        if [ "$booted" = "$built" ]; then
+        ${nixos-rebuild} boot ${toString (cfg.flags ++ upgradeFlag)}
+        booted="$(${readlink} /run/booted-system/{initrd,kernel,kernel-modules})"
+        built="$(${readlink} /nix/var/nix/profiles/system/{initrd,kernel,kernel-modules})"
+
+        ${optionalString (cfg.rebootWindow != null) ''
+          current_time="$(${date} +%H:%M)"
+
+          lower="${cfg.rebootWindow.lower}"
+          upper="${cfg.rebootWindow.upper}"
+
+          if [[ "''${lower}" < "''${upper}" ]]; then
+            if [[ "''${current_time}" > "''${lower}" ]] && \
+               [[ "''${current_time}" < "''${upper}" ]]; then
+              do_reboot="true"
+            else
+              do_reboot="false"
+            fi
+          else
+            # lower > upper, so we are crossing midnight (e.g. lower=23h, upper=6h)
+            # we want to reboot if cur > 23h or cur < 6h
+            if [[ "''${current_time}" < "''${upper}" ]] || \
+               [[ "''${current_time}" > "''${lower}" ]]; then
+              do_reboot="true"
+            else
+              do_reboot="false"
+            fi
+          fi
+        ''}
+
+        if [ "''${booted}" = "''${built}" ]; then
           ${nixos-rebuild} switch ${toString cfg.flags}
+        ${optionalString (cfg.rebootWindow != null) ''
+          elif [ "''${do_reboot}" != true ]; then
+            echo "Outside of configured reboot window, skipping."
+        ''}
         else
-          /run/current-system/sw/bin/shutdown -r +1
+          ${shutdown} -r +1
         fi
       '' else ''
-        ${nixos-rebuild} switch ${toString cfg.flags}
+        ${nixos-rebuild} switch ${toString (cfg.flags ++ upgradeFlag)}
       '';
 
       startAt = cfg.dates;
@@ -167,3 +225,4 @@ in {
   };
 
 }
+