Fixing NVMe Overheating After Suspend

1
{ config, pkgs, lib, ... }:
2

3
let
4
  cfg = config.hardware.SsdApstFix;
5
in
6
{
7
  options.hardware.SsdApstFix = {
8
    serial = lib.mkOption {
9
      type = lib.types.nullOr lib.types.str;
10
      default = null;
11
      description = "Serial number of the NVMe drive to apply the PS4 resume fix to.";
12
    };
13
  };
14

15
  config = lib.mkIf (cfg.serial != null) {
16
    environment.systemPackages = [ pkgs.nvme-cli ];
17

18
    # NixOS provides a really nice way to write post-resume hooks.
19
    powerManagement.resumeCommands =
20
      let
21
        nvmeFix = pkgs.writeShellScript "ssd-force-ps4-on-resume" ''
22
          set -euo pipefail
23

24
          export PATH=${lib.makeBinPath [ pkgs.coreutils pkgs.gawk pkgs.nvme-cli pkgs.systemd ]}
25

26
          TAG="nvme-ps4-fix"
27
          log_info() { echo "$*" | systemd-cat -t "$TAG" -p info; }
28
          log_err()  { echo "$*" | systemd-cat -t "$TAG" -p err; }
29

30
          SERIAL=${lib.escapeShellArg cfg.serial}
31
          TARGET_PS="4"
32

33
          get_ps() {
34
            nvme get-feature "$1" -f 0x02 -H 2>/dev/null | awk '/Power State/{print $NF; exit}'
35
          }
36

37
          # Wait up to ~5s for NVMe sysfs to settle after resume
38
          for outer in $(seq 1 25); do
39
            for c in /sys/class/nvme/nvme*; do
40
              [ -r "$c/serial" ] || continue
41

42
              # NVMe strings can be space padded
43
              s="$(tr -d ' \n' < "$c/serial")"
44
              [ "$s" = "$SERIAL" ] || continue
45

46
              dev="/dev/$(basename "$c")"
47

48
              cur_ps="$(get_ps "$dev" || true)"
49
              [ -n "$cur_ps" ] || cur_ps="unknown"
50

51
              if [ "$cur_ps" = "$TARGET_PS" ]; then
52
                log_info "matched serial=$SERIAL dev=$dev; already in PS$TARGET_PS"
53
                exit 0
54
              fi
55

56
              new_ps="unknown"
57

58
              # Try to set PS4 (retry; resume can be racy)
59
              for i in $(seq 1 10); do
60
                nvme set-feature "$dev" -f 0x02 -V 0x4 >/dev/null 2>&1 || true
61

62
                new_ps="$(get_ps "$dev" || true)"
63
                [ -n "$new_ps" ] || new_ps="unknown"
64

65
                if [ "$new_ps" = "$TARGET_PS" ]; then
66
                  log_info "matched serial=$SERIAL dev=$dev; PS $cur_ps -> $new_ps"
67
                  exit 0
68
                fi
69

70
                sleep 0.2
71
              done
72

73
              log_err "matched serial=$SERIAL dev=$dev; failed to confirm PS4 (was PS=$cur_ps, now PS=$new_ps)"
74
              exit 0
75
            done
76

77
            sleep 0.2
78
          done
79

80
          log_err "no NVMe controller found with serial=$SERIAL"
81
          exit 0
82
        '';
83
      in
84
      ''
85
        ${nvmeFix}
86
      '';
87
  };
88
}
89

90
# Set hardware.SsdApstFix.serial = "xyz"; after importing it!