From bd3e9c3d051f77338b38ea44cc28cfc465923f41 Mon Sep 17 00:00:00 2001
From: Graham Christensen <graham@grahamc.com>
Date: Tue, 8 Mar 2022 09:28:41 -0500
Subject: [PATCH] nixos/zfs: better support auto-expanding partitioned disks

---
 nixos/modules/tasks/filesystems/zfs.nix | 35 +++----------------------
 1 file changed, 3 insertions(+), 32 deletions(-)

diff --git a/nixos/modules/tasks/filesystems/zfs.nix b/nixos/modules/tasks/filesystems/zfs.nix
index 3bc0dedec00..fbfc61177d3 100644
--- a/nixos/modules/tasks/filesystems/zfs.nix
+++ b/nixos/modules/tasks/filesystems/zfs.nix
@@ -642,41 +642,14 @@ in
         };
 
         scriptArgs = "%i";
-        path = [ pkgs.gawk cfgZfs.package ];
+        path = [ cfgZfs.package ];
 
-        # ZFS has no way of enumerating just devices in a pool in a way
-        # that 'zpool online -e' supports. Thus, we've implemented a
-        # bit of a strange approach of highlighting just devices.
-        # See: https://github.com/openzfs/zfs/issues/12505
-        script = let
-          # This UUID has been chosen at random and is to provide a
-          # collision-proof, predictable token to search for
-          magicIdentifier = "NIXOS-ZFS-ZPOOL-DEVICE-IDENTIFIER-37108bec-aff6-4b58-9e5e-53c7c9766f05";
-          zpoolScripts = pkgs.writeShellScriptBin "device-highlighter" ''
-            echo "${magicIdentifier}"
-          '';
-        in ''
+        script =  ''
           pool=$1
 
           echo "Expanding all devices for $pool."
 
-          # Put our device-highlighter script it to the PATH
-          export ZPOOL_SCRIPTS_PATH=${zpoolScripts}/bin
-
-          # Enable running our precisely specified zpool script as root
-          export ZPOOL_SCRIPTS_AS_ROOT=1
-
-          devices() (
-            zpool status -c device-highlighter "$pool" \
-             | awk '($2 == "ONLINE" && $6 == "${magicIdentifier}") { print $1; }'
-          )
-
-          for device in $(devices); do
-            echo "Attempting to expand $device of $pool..."
-            if ! zpool online -e "$pool" "$device"; then
-              echo "Failed to expand '$device' of '$pool'."
-            fi
-          done
+          ${pkgs.zpool-auto-expand-partitions}/bin/zpool_part_disks --automatically-grow "$pool"
         '';
       };
 
@@ -701,8 +674,6 @@ in
             RemainAfterExit = true;
           };
 
-          path = [ pkgs.gawk cfgZfs.package ];
-
           script = ''
             for pool in ${poolListProvider}; do
               systemctl start --no-block "zpool-expand@$pool"