From cd3ffbdfe6869c35c483bc7cec348e8d5365c2b3 Mon Sep 17 00:00:00 2001
From: Nick Cao <nickcao@nichi.co>
Date: Mon, 20 Mar 2023 22:50:11 +0800
Subject: [PATCH] nixos/zram: add writebackDevice option and corresponding test

---
 nixos/modules/config/zram.nix  | 20 ++++++++++++++++++++
 nixos/tests/zram-generator.nix | 34 ++++++++++++++++++++++++++--------
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/nixos/modules/config/zram.nix b/nixos/modules/config/zram.nix
index 4df646cf27966..991387ea9b2bd 100644
--- a/nixos/modules/config/zram.nix
+++ b/nixos/modules/config/zram.nix
@@ -82,12 +82,30 @@ in
           {command}`cat /sys/class/block/zram*/comp_algorithm`
         '';
       };
+
+      writebackDevice = lib.mkOption {
+        default = null;
+        example = "/dev/zvol/tarta-zoot/swap-writeback";
+        type = lib.types.nullOr lib.types.path;
+        description = lib.mdDoc ''
+          Write incompressible pages to this device,
+          as there's no gain from keeping them in RAM.
+        '';
+      };
     };
 
   };
 
   config = lib.mkIf cfg.enable {
 
+    assertions = [
+      {
+        assertion = cfg.writebackDevice == null || cfg.swapDevices <= 1;
+        message = "A single writeback device cannot be shared among multiple zram devices";
+      }
+    ];
+
+
     system.requiredKernelConfig = with config.lib.kernelConfig; [
       (isModule "ZRAM")
     ];
@@ -112,6 +130,8 @@ in
                 zram-size = if cfg.memoryMax != null then "min(${size}, ${toString cfg.memoryMax} / 1024 / 1024)" else size;
                 compression-algorithm = cfg.algorithm;
                 swap-priority = cfg.priority;
+              } // lib.optionalAttrs (cfg.writebackDevice != null) {
+                writeback-device = cfg.writebackDevice;
               };
           })
           devices));
diff --git a/nixos/tests/zram-generator.nix b/nixos/tests/zram-generator.nix
index 3407361d2824f..2be7bd2e05b1d 100644
--- a/nixos/tests/zram-generator.nix
+++ b/nixos/tests/zram-generator.nix
@@ -1,18 +1,36 @@
 import ./make-test-python.nix {
   name = "zram-generator";
 
-  nodes.machine = { ... }: {
-    zramSwap = {
-      enable = true;
-      priority = 10;
-      algorithm = "lz4";
-      swapDevices = 2;
-      memoryPercent = 30;
-      memoryMax = 10 * 1024 * 1024;
+  nodes = {
+    single = { ... }: {
+      virtualisation = {
+        emptyDiskImages = [ 512 ];
+      };
+      zramSwap = {
+        enable = true;
+        priority = 10;
+        algorithm = "lz4";
+        swapDevices = 1;
+        memoryPercent = 30;
+        memoryMax = 10 * 1024 * 1024;
+        writebackDevice = "/dev/vdb";
+      };
+    };
+    machine = { ... }: {
+      zramSwap = {
+        enable = true;
+        priority = 10;
+        algorithm = "lz4";
+        swapDevices = 2;
+        memoryPercent = 30;
+        memoryMax = 10 * 1024 * 1024;
+      };
     };
   };
 
   testScript = ''
+    single.wait_for_unit("systemd-zram-setup@zram0.service")
+
     machine.wait_for_unit("systemd-zram-setup@zram0.service")
     machine.wait_for_unit("systemd-zram-setup@zram1.service")
     zram = machine.succeed("zramctl --noheadings --raw")