dotfiles/nixos/modules/zfs-health.nix

{
  config,
  pkgs,
  lib,
  ...
}: let
  cfg = config.zfs.health;
  pool = {
    config,
    name,
    ...
  }: {
    options = {
      scrub = lib.mkOption {
        type = with lib.types; nullOr str;
      };

      trim = {
        secure = lib.mkOption {
          type = lib.types.bool;
          default = false;
        };

        rate = lib.mkOption {
          type = with lib.types; nullOr int;
          default = null;
        };

        schedule = lib.mkOption {
          type = with lib.types; nullOr str;
        };
      };
    };
  };

  zpool = lib.getExe' cfg.package "zpool";

  trimPools = lib.filterAttrs (_: poolConfig: poolConfig.trim != null) cfg.pools;
in {
  options.zfs.health = {
    enable = lib.mkEnableOption "Enable ZFS health assurance module.";

    package = lib.mkOption {
      type = lib.types.package;
      default = config.boot.zfs.package;
    };

    pools = lib.mkOption {
      type = with lib.types; attrsOf (submodule pool);
      default = [];
    };
  };

  config = lib.mkIf cfg.enable {
    systemd.services = lib.foldl (acc: attr: acc // attr) {} [
      (lib.flip lib.mapAttrs' trimPools
        (
          pool: poolConfig:
            lib.nameValuePair "zfs-health-trim-${pool}-resume" {
              after = ["systemd-suspend.service"];
              wantedBy = ["suspend.target"];

              restartIfChanged = false;

              script = ''
                set -eo pipefail

                if -e /var/run/zfs-health-trim-${pool}-suspended ; then
                  systemctl start "zfs-healt-trim-${pool}.service"
                  rm /var/run/zfs-health-trim-${pool}-suspended
                fi
              '';

              serviceConfig = {
                Type = "oneshot";
              };
            }
        ))
      (lib.flip lib.mapAttrs' trimPools
        (
          pool: poolConfig:
            lib.nameValuePair "zfs-health-trim-${pool}-suspend" {
              before = ["systemd-suspend.service"];
              wantedBy = ["suspend.target"];

              restartIfChanged = false;

              script = ''
                set -eo pipefail

                if systemctl status "zfs-healt-trim-${pool}.service" ; then
                  touch /var/run/zfs-health-trim-${pool}-suspended
                fi
              '';

              serviceConfig = {
                Type = "oneshot";
              };
            }
        ))
      (lib.flip lib.mapAttrs' trimPools
        (pool: poolConfig:
          lib.nameValuePair "zfs-health-trim-${pool}" {
            after = ["multi-user.target"];

            restartIfChanged = false;

            preStop = ''
              set -eo pipefail

              if ${zpool} status ${pool} | grep -q trimming ; then
                ${zpool} trim --suspend ${pool}

                if ! [ -z $MAINPID ] ; then
                  while kill -0 $MAINPID 2> /dev/null ; do
                    sleep 10
                  done
                fi
              fi
            '';

            script = ''
              if ${zpool} status ${pool} | grep -q trimming ; then
                echo "Trim already in progress, polling for completion"
                while ${zpool} status ${pool} | grep trimming >/dev/null ; do
                  sleep 60
                done
              else
                ${zpool} trim ${pool} \
                  --wait \
                  ${lib.optionalString (poolConfig.trim.rate != null) "--rate ${toString poolConfig.trim.rate}"} \
                  ${lib.optionalString poolConfig.trim.secure "--secure"}
              fi

              echo "Trim done!"
              ${zpool} status ${pool}
            '';

            serviceConfig = {
              Type = "oneshot";
            };
          }))
    ];

    systemd.timers = lib.pipe cfg.pools [
      (lib.filterAttrs (_: poolConfig: poolConfig.trim != null))
      (lib.mapAttrs'
        (pool: poolConfig:
          lib.nameValuePair "zfs-health-trim-${pool}" {
            after = ["multi-user.target"];

            timerConfig = {
              OnCalendar = poolConfig.trim.schedule;
              Persistent = true;
            };
          }))
    ];
  };
}