From cb9cf2ea0a2e30f27f8ce5a1d5a6a855a8b91438 Mon Sep 17 00:00:00 2001
From: Ameer Hamza <ahamza@ixsystems.com>
Date: Tue, 14 May 2024 19:40:30 +0500
Subject: [PATCH] nvme-of: Remove/Add disk on discovery change event

---
 .../nvmf-connect@.service.d/override.conf     |  3 ++
 src/freenas/usr/local/bin/nvmf-connect.sh     | 50 +++++++++++++++++++
 src/freenas/usr/local/bin/nvmf-wrapper.sh     | 29 +++++++++++
 3 files changed, 82 insertions(+)
 create mode 100644 src/freenas/etc/systemd/system/nvmf-connect@.service.d/override.conf
 create mode 100755 src/freenas/usr/local/bin/nvmf-connect.sh
 create mode 100755 src/freenas/usr/local/bin/nvmf-wrapper.sh

diff --git a/src/freenas/etc/systemd/system/nvmf-connect@.service.d/override.conf b/src/freenas/etc/systemd/system/nvmf-connect@.service.d/override.conf
new file mode 100644
index 000000000000..e216fa3ee6e3
--- /dev/null
+++ b/src/freenas/etc/systemd/system/nvmf-connect@.service.d/override.conf
@@ -0,0 +1,3 @@
+[Service]
+ExecStart=
+ExecStart=/bin/sh -c "/usr/local/bin/nvmf-wrapper.sh `/bin/echo -e '${CONNECT_ARGS}'`"
diff --git a/src/freenas/usr/local/bin/nvmf-connect.sh b/src/freenas/usr/local/bin/nvmf-connect.sh
new file mode 100755
index 000000000000..ff91555ae62f
--- /dev/null
+++ b/src/freenas/usr/local/bin/nvmf-connect.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+
+RECONNECT_TIMEOUT_SEC=5
+discovery_dev=$(echo "$@" | grep -oP '(?<=--device=)\w+')
+discovery_addr=$(</sys/class/nvme/${discovery_dev}/address)
+queue="/var/run/${discovery_dev}-fabric-queue"
+
+while [ -s $queue ]; do
+  truncate -s 0 $queue
+  local_nqns=""
+  disconnected=0
+  discovery_nqns=$(nvme discover $@ | grep subnqn: | awk '{gsub(/^subnqn: */,"")}1')
+
+  for nqn_f in /sys/class/nvme-fabrics/ctl/nvme*/subsysnqn; do
+    dev=$(basename "$(dirname "$nqn_f")")
+    nqn=$(<"$nqn_f")
+    addr=$(</sys/class/nvme/${dev}/address)
+
+    if [[ "$dev" == "$discovery_dev" || "$addr" != "$discovery_addr" ]]; then
+      continue
+    fi
+
+    if [[ "$discovery_nqns" != *"$nqn"* ]]; then
+      nvme disconnect -d $dev
+      disconnected=1
+    else
+      local_nqns+="$nqn"$'\n'
+    fi
+  done
+
+  if [ $disconnected -eq 0 ]; then
+    counter=0
+    local_nqns=$(echo -e "$local_nqns" | tr -s '\n' | sort)
+
+    while [ $counter -lt $RECONNECT_TIMEOUT_SEC ]; do
+      connect_nqns=$(nvme discover "$@" | grep subnqn: | awk '{gsub(/^subnqn: */,"")}1' \
+          | tr -s '\n' | sort)
+
+      if [ "$local_nqns" != "$connect_nqns" ]; then
+        break
+      fi
+
+      ((counter++))
+      sleep 1
+    done
+  fi
+
+  connect_all=$(nvme connect-all $@)
+done
+echo "EXIT" > $queue
diff --git a/src/freenas/usr/local/bin/nvmf-wrapper.sh b/src/freenas/usr/local/bin/nvmf-wrapper.sh
new file mode 100755
index 000000000000..4b4ae990576f
--- /dev/null
+++ b/src/freenas/usr/local/bin/nvmf-wrapper.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+
+if [ "$#" -eq 0 ]; then
+  echo "Error: No arguments provided."
+  exit 1
+fi
+
+discovery_nqn="nqn.2014-08.org.nvmexpress.discovery"
+discovery_dev=$(echo "$@" | grep -oP '(?<=--device=)\w+')
+
+if [[ $(</sys/class/nvme/${discovery_dev}/subsysnqn) != $discovery_nqn ]]; then
+  echo "Error: Event not on discovery controller"
+  exit 1
+fi
+
+queue="/var/run/${discovery_dev}-fabric-queue"
+echo "QUEUE" > $queue
+systemd-run --unit=${discovery_dev}-of.service /bin/bash -c "/usr/local/bin/nvmf-connect.sh $*"
+
+# If service is still active and queue reads EXIT, we have a race. Wait for small interval since
+# we already know that nvmf-connect.sh is about to exit.
+if [ $? != 0 ] && [ "EXIT" = $(<$queue) ]; then
+  sleep 0.1
+  systemd-run --unit=${discovery_dev}-of.service /bin/bash -c "/usr/local/bin/nvmf-connect.sh $*"
+
+  if [ $? != 0 ]; then
+    echo "Error: ${discovery_dev}-of.service failed to start after waiting for 100 milliseconds"
+  fi
+fi