fix(disk/lsblk): support older lsblk without JSON mode, using --pairs

pkg/disk/lsblk.go

@@ -24,35 +24,109 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"regexp"
 	"sort"
+	"strconv"
 	"strings"
 
 	"github.com/dustin/go-humanize"
 	"github.com/olekukonko/tablewriter"
 	"sigs.k8s.io/yaml"
 
+	"github.com/leptonai/gpud/log"
 	"github.com/leptonai/gpud/pkg/file"
 	"github.com/leptonai/gpud/pkg/process"
 )
 
 const (
+	lsblkVersionFlags = "--version"
 	// lsblkFlags adds device name, if add empty string - command will print info about all devices
-	lsblkFlags = "--paths --json --bytes --fs --output NAME,TYPE,SIZE,ROTA,SERIAL,WWN,VENDOR,MODEL,REV,MOUNTPOINT,FSTYPE,PARTUUID"
+	lsblkFlags = "--paths --bytes --fs --output NAME,TYPE,SIZE,ROTA,SERIAL,WWN,VENDOR,MODEL,REV,MOUNTPOINT,FSTYPE,PARTUUID"
+	// lsblkJsonFlag lsblk from version 2.37 support json response
+	lsblkJsonFlag = "--json"
+	// lsblkMinSupportJsonVersion lsblk from version 2.37 support json response
+	// https://github.com/util-linux/util-linux/blob/stable/v2.27/misc-utils/lsblk.c#L1626
+	lsblkMinSupportJsonVersion = 2.37
+	// lsblkPairsFlag lsblk lower than 2.37 only support raw and pairs response
+	lsblkPairsFlag = "--pairs"
 	// outputKey is the key to find block devices in lsblk json output
 	outputKey = "blockdevices"
 )
 
+var lsblkVersionRegPattern = regexp.MustCompile(`\d+\.\d+`)
+
+// decideLsblkFlagAndParserFromVersion decides the lsblk command flags based on the "lsblk --version" output
+func decideLsblkFlagAndParserFromVersion(verOutput string) (string, func([]byte, ...OpOption) (BlockDevices, error), error) {
+	matches := lsblkVersionRegPattern.FindString(verOutput)
+	if matches != "" {
+		if versionF, parseErr := strconv.ParseFloat(matches, 64); parseErr == nil {
+			if versionF >= lsblkMinSupportJsonVersion {
+				return lsblkFlags + " " + lsblkJsonFlag, ParseJSON, nil
+			}
+
+			return lsblkFlags + " " + lsblkPairsFlag, ParsePairs, nil
+		}
+	}
+
+	return "", nil, errors.New("failed to parse 'lsblk --version' output")
+}
+
+func decideLsblkFlag(ctx context.Context) (string, func([]byte, ...OpOption) (BlockDevices, error), error) {
+	lsblkVersion, err := file.LocateExecutable("lsblk")
+	if err != nil {
+		return "", nil, err
+	}
+
+	p, err := process.New(
+		process.WithCommand(lsblkVersion+" "+lsblkVersionFlags),
+		process.WithRunAsBashScript(),
+	)
+	if err != nil {
+		return "", nil, err
+	}
+
+	if err := p.Start(ctx); err != nil {
+		return "", nil, err
+	}
+
+	lines := make([]string, 0)
+	if err := process.Read(
+		ctx,
+		p,
+		process.WithReadStdout(),
+		process.WithReadStderr(),
+		process.WithProcessLine(func(line string) {
+			lines = append(lines, line)
+		}),
+		process.WithWaitForCmd(),
+	); err != nil {
+		return "", nil, fmt.Errorf("failed to check lsblk version: %w", err)
+	}
+
+	line := strings.Join(lines, "\n")
+	line = strings.TrimSpace(line)
+
+	return decideLsblkFlagAndParserFromVersion(line)
+}
+
 // GetBlockDevices run os lsblk command for device and construct BlockDevice struct based on output
 // Receives device path. If device is empty string, info about all devices will be collected
 // Returns slice of BlockDevice structs or error if something went wrong
 func GetBlockDevices(ctx context.Context, opts ...OpOption) (BlockDevices, error) {
 	lsblkPath, err := file.LocateExecutable("lsblk")
 	if err != nil {
-		return nil, nil
+		return nil, err
+	}
+
+	// pre-check lsblk version
+	flags, parseFunc, checkErr := decideLsblkFlag(ctx)
+	if checkErr != nil {
+		log.Logger.Warnw("failed to decide lsblk flag and parser -- falling back to latest version", "error", checkErr)
+		flags, parseFunc = lsblkFlags+" "+lsblkJsonFlag, ParseJSON
 	}
 
 	p, err := process.New(
-		process.WithCommand(lsblkPath+" "+lsblkFlags),
+		process.WithCommand(lsblkPath+" "+flags),
 		process.WithRunAsBashScript(),
 	)
 	if err != nil {
@@ -77,10 +151,10 @@ func GetBlockDevices(ctx context.Context, opts ...OpOption) (BlockDevices, error
 		return nil, fmt.Errorf("failed to read lsblk output: %w\n\noutput:\n%s", err, strings.Join(lines, "\n"))
 	}
 
-	return Parse([]byte(strings.Join(lines, "\n")), opts...)
+	return parseFunc([]byte(strings.Join(lines, "\n")), opts...)
 }
 
-func Parse(b []byte, opts ...OpOption) (BlockDevices, error) {
+func ParseJSON(b []byte, opts ...OpOption) (BlockDevices, error) {
 	if len(b) == 0 {
 		return nil, errors.New("empty input provided to Parse")
 	}
@@ -135,6 +209,143 @@ func Parse(b []byte, opts ...OpOption) (BlockDevices, error) {
 	return devs, nil
 }
 
+func ParsePairs(b []byte, opts ...OpOption) (BlockDevices, error) {
+	if len(b) == 0 {
+		return nil, errors.New("empty input provided to ParsePairs")
+	}
+
+	devs := make(BlockDevices, 0)
+
+	// parse each line
+	lines := strings.Split(string(b), "\n")
+	for _, line := range lines {
+		// skip empty line
+		if len(line) == 0 {
+			continue
+		}
+
+		// parse each row then return BlockDevice
+		disk, err := parseLineToDisk(line)
+		if err != nil {
+			return nil, err
+		}
+
+		// parse each block then add blocks slice
+		devs = append(devs, disk)
+	}
+
+	// build disk hierarchy
+	devs = buildDiskHierarchy(devs)
+	if len(devs) == 0 {
+		return nil, errors.New("build disk hierarchy failed")
+	}
+
+	// to JSON bytes
+	jsonData, err := json.MarshalIndent(struct {
+		BlockDevices BlockDevices `json:"blockdevices"`
+	}{BlockDevices: devs}, "", "  ")
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal lsblk-blockdevices json mode")
+	}
+
+	return ParseJSON(jsonData, opts...)
+}
+
+func parseLineToDisk(line string) (BlockDevice, error) {
+	disk := BlockDevice{}
+	parts := strings.Fields(line)
+
+	for _, part := range parts {
+		kv := strings.Split(part, "=")
+		if len(kv) != 2 {
+			continue
+		}
+		key, value := kv[0], strings.Trim(kv[1], `"`)
+		switch key {
+		case "NAME":
+			disk.Name = value
+		case "TYPE":
+			disk.Type = value
+		case "SIZE":
+			disk.Size = toCustomInt64(value)
+		case "ROTA":
+			disk.Rota = toCustomBool(value)
+		case "SERIAL":
+			disk.Serial = value
+		case "WWN":
+			disk.WWN = value
+		case "VENDOR":
+			disk.Vendor = value
+		case "MODEL":
+			disk.Model = value
+		case "REV":
+			disk.Rev = value
+		case "MOUNTPOINT":
+			disk.MountPoint = value
+		case "FSTYPE":
+			disk.FSType = value
+		case "PARTUUID":
+			disk.PartUUID = value
+		case "PKNAME":
+			disk.PKName = value
+		}
+	}
+
+	return disk, nil
+}
+
+func buildDiskHierarchy(disks BlockDevices) (finalDisks BlockDevices) {
+	// Recursive function to nest child disks into their parent disks
+	var recursiveAdd func(disk BlockDevice, disks *BlockDevices)
+
+	// Implementation of the recursive nesting function
+	recursiveAdd = func(disk BlockDevice, disks *BlockDevices) {
+		// Find the parent disk of the current disk and recursively nest
+		for i := range *disks {
+			if (*disks)[i].Name == disk.PKName {
+				// Found the parent disk, add the current disk to the parent's Children
+				(*disks)[i].Children = append((*disks)[i].Children, disk)
+				return
+			}
+
+			// If the current disk has children, continue recursively
+			recursiveAdd(disk, (*BlockDevices)(&(*disks)[i].Children))
+		}
+	}
+
+	// Add disks that don't have a parent disk to finalDisks
+	for i := range disks {
+		if disks[i].PKName == "" {
+			finalDisks = append(finalDisks, disks[i])
+		}
+	}
+
+	// Perform recursive nesting for each disk
+	for i := range disks {
+		if disks[i].PKName != "" {
+			recursiveAdd(disks[i], &finalDisks)
+		}
+	}
+
+	return finalDisks
+}
+
+func toCustomInt64(value string) CustomInt64 {
+	n, err := strconv.ParseInt(value, 10, 64)
+	if err != nil {
+		return CustomInt64{}
+	}
+	return CustomInt64{n}
+}
+
+func toCustomBool(value string) CustomBool {
+	n, err := strconv.ParseBool(value)
+	if err != nil {
+		return CustomBool{}
+	}
+	return CustomBool{n}
+}
+
 type BlockDevices []BlockDevice
 
 func (blks BlockDevices) JSON() ([]byte, error) {
@@ -199,6 +410,7 @@ type BlockDevice struct {
 	MountPoint       string        `json:"mountpoint,omitempty"`
 	FSType           string        `json:"fstype,omitempty"`
 	PartUUID         string        `json:"partuuid,omitempty"`
+	PKName           string        `json:"-"`
 	Children         []BlockDevice `json:"children,omitempty"`
 }
 

pkg/disk/lsblk_test.go

@@ -20,6 +20,7 @@ package disk
 
 import (
 	"os"
+	"reflect"
 	"testing"
 
 	"github.com/dustin/go-humanize"
@@ -34,13 +35,13 @@ func TestParse(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		blks, err := Parse(dat)
+		blks, err := ParseJSON(dat)
 		if err != nil {
 			t.Fatal(err)
 		}
 		blks.RenderTable(os.Stdout)
 
-		blks, err = Parse(dat, WithDeviceType(func(deviceType string) bool {
+		blks, err = ParseJSON(dat, WithDeviceType(func(deviceType string) bool {
 			return deviceType == "disk"
 		}))
 		if err != nil {
@@ -51,3 +52,45 @@ func TestParse(t *testing.T) {
 		t.Logf("Total bytes: %s", humanize.Bytes(totalBytes))
 	}
 }
+
+func TestParsePairs(t *testing.T) {
+	t.Parallel()
+
+	for _, f := range []string{"lsblk.3.txt"} {
+		dat, err := os.ReadFile("testdata/" + f)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		blks, err := ParsePairs(dat, WithDeviceType(func(deviceType string) bool {
+			return deviceType == "disk"
+		}))
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		blks.RenderTable(os.Stdout)
+		totalBytes := blks.GetTotalBytes()
+		t.Logf("Total bytes: %s", humanize.Bytes(totalBytes))
+	}
+}
+
+func TestCheckVersion(t *testing.T) {
+	t.Parallel()
+
+	expecteds := []string{
+		"--paths --bytes --fs --output NAME,TYPE,SIZE,ROTA,SERIAL,WWN,VENDOR,MODEL,REV,MOUNTPOINT,FSTYPE,PARTUUID --pairs",
+		"--paths --bytes --fs --output NAME,TYPE,SIZE,ROTA,SERIAL,WWN,VENDOR,MODEL,REV,MOUNTPOINT,FSTYPE,PARTUUID --json",
+		"--paths --bytes --fs --output NAME,TYPE,SIZE,ROTA,SERIAL,WWN,VENDOR,MODEL,REV,MOUNTPOINT,FSTYPE,PARTUUID --json",
+	}
+
+	for i, s := range []string{"lsblk，来自 util-linux 2.23.2", "lsblk from util-linux 2.37.2", "lsblk from util-linux 2.37.4"} {
+		lsblkCmd, _, err := decideLsblkFlagAndParserFromVersion(s)
+		if err != nil {
+			t.Errorf("Expected %v, got %v", expecteds[i], lsblkCmd)
+		}
+		if !reflect.DeepEqual(lsblkCmd, expecteds[i]) {
+			t.Errorf("Expected %v, got %v", expecteds[i], lsblkCmd)
+		}
+	}
+}

pkg/disk/testdata/lsblk.3.txt

@@ -0,0 +1,8 @@
+NAME="/dev/sda" TYPE="disk" SIZE="959656755200" ROTA="0" SERIAL="600062b20b6056002b4c9a038d0f9a62" WWN="0x600062b20b605600" VENDOR="AVAGO   " MODEL="INSPUR          " REV="5.22" MOUNTPOINT="" FSTYPE="" PARTUUID="" PKNAME=""
+NAME="/dev/sda1" TYPE="part" SIZE="209715200" ROTA="0" SERIAL="" WWN="0x600062b20b605600" VENDOR="" MODEL="" REV="" MOUNTPOINT="/boot/efi" FSTYPE="vfat" PARTUUID="57928472-752e-42d0-a910-d293e87d21b8" PKNAME="/dev/sda"
+NAME="/dev/sda2" TYPE="part" SIZE="1073741824" ROTA="0" SERIAL="" WWN="0x600062b20b605600" VENDOR="" MODEL="" REV="" MOUNTPOINT="/boot" FSTYPE="xfs" PARTUUID="2cb1faf9-e22c-4fda-b8e0-761093c0f89c" PKNAME="/dev/sda"
+NAME="/dev/sda3" TYPE="part" SIZE="958371201024" ROTA="0" SERIAL="" WWN="0x600062b20b605600" VENDOR="" MODEL="" REV="" MOUNTPOINT="" FSTYPE="LVM2_member" PARTUUID="399fae2e-c11a-4f06-9257-8274b7b40014" PKNAME="/dev/sda"
+NAME="/dev/mapper/centos-root" TYPE="lvm" SIZE="924009365504" ROTA="0" SERIAL="" WWN="" VENDOR="" MODEL="" REV="" MOUNTPOINT="/" FSTYPE="xfs" PARTUUID="" PKNAME="/dev/sda3"
+NAME="/dev/mapper/centos-swap" TYPE="lvm" SIZE="34359738368" ROTA="0" SERIAL="" WWN="" VENDOR="" MODEL="" REV="" MOUNTPOINT="" FSTYPE="swap" PARTUUID="" PKNAME="/dev/sda3"
+NAME="/dev/nvme0n1" TYPE="disk" SIZE="3200631791616" ROTA="0" SERIAL="PHLN231204BN3P2FGN" WWN="eui.01000000010000005cd2e44f8f7e5551" VENDOR="" MODEL="INTEL SSDPE2KE032T8                     " REV="" MOUNTPOINT="/data1" FSTYPE="xfs" PARTUUID="" PKNAME=""
+NAME="/dev/nvme1n1" TYPE="disk" SIZE="3200631791616" ROTA="0" SERIAL="PHLN2312047T3P2FGN" WWN="eui.01000000010000005cd2e4cc8e7e5551" VENDOR="" MODEL="INTEL SSDPE2KE032T8                     " REV="" MOUNTPOINT="/data2" FSTYPE="xfs" PARTUUID="" PKNAME=""