From 5f4f64964d27f2a8842aedf306e22d2b830f870d Mon Sep 17 00:00:00 2001 From: Gyuho Lee Date: Tue, 31 Dec 2024 14:57:22 +0800 Subject: [PATCH] use parse pairs for older lsblk Signed-off-by: Gyuho Lee --- pkg/disk/lsblk.go | 54 +++++++++++++++++++++++------------------- pkg/disk/lsblk_test.go | 13 ++++++---- 2 files changed, 38 insertions(+), 29 deletions(-) diff --git a/pkg/disk/lsblk.go b/pkg/disk/lsblk.go index f5697f69..92ed76ed 100644 --- a/pkg/disk/lsblk.go +++ b/pkg/disk/lsblk.go @@ -33,6 +33,7 @@ import ( "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" ) @@ -54,25 +55,26 @@ const ( var lsblkVersionRegPattern = regexp.MustCompile(`\d+\.\d+`) -func lsblkVersionCheck(line string) string { - matches := lsblkVersionRegPattern.FindString(line) +// 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 - } else { - return lsblkFlags + " " + lsblkPairsFlag + return lsblkFlags + " " + lsblkJsonFlag, ParseJSON, nil } + + return lsblkFlags + " " + lsblkPairsFlag, ParsePairs, nil } } - return "" + return "", nil, errors.New("failed to parse 'lsblk --version' output") } -func CheckLsblk(ctx context.Context) (string, error) { +func decideLsblkFlag(ctx context.Context) (string, func([]byte, ...OpOption) (BlockDevices, error), error) { lsblkVersion, err := file.LocateExecutable("lsblk") if err != nil { - return "", err + return "", nil, err } p, err := process.New( @@ -80,47 +82,51 @@ func CheckLsblk(ctx context.Context) (string, error) { process.WithRunAsBashScript(), ) if err != nil { - return "", err + return "", nil, err } if err := p.Start(ctx); err != nil { - return "", err + return "", nil, err } - var lsblkCmd string + lines := make([]string, 0) if err := process.Read( ctx, p, process.WithReadStdout(), process.WithReadStderr(), process.WithProcessLine(func(line string) { - lsblkCmd = lsblkVersionCheck(line) + lines = append(lines, line) }), process.WithWaitForCmd(), ); err != nil { - return "", fmt.Errorf("failed to check lsblk version: %w", err) + return "", nil, fmt.Errorf("failed to check lsblk version: %w", err) } - return lsblkCmd, nil + 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) { - // pre-check lsblk version - lsblkExecuteFlags, checkErr := CheckLsblk(ctx) - if checkErr != nil { - return nil, checkErr - } - lsblkPath, err := file.LocateExecutable("lsblk") if err != 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+" "+lsblkExecuteFlags), + process.WithCommand(lsblkPath+" "+flags), process.WithRunAsBashScript(), ) if err != nil { @@ -145,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") } @@ -242,7 +248,7 @@ func ParsePairs(b []byte, opts ...OpOption) (BlockDevices, error) { return nil, fmt.Errorf("failed to marshal lsblk-blockdevices json mode") } - return Parse(jsonData, opts...) + return ParseJSON(jsonData, opts...) } func parseLineToDisk(line string) (BlockDevice, error) { diff --git a/pkg/disk/lsblk_test.go b/pkg/disk/lsblk_test.go index dc6239b0..4ec1ac37 100644 --- a/pkg/disk/lsblk_test.go +++ b/pkg/disk/lsblk_test.go @@ -35,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 { @@ -81,11 +81,14 @@ func TestCheckVersion(t *testing.T) { 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.4"} { - var lsblkCmd string - lsblkCmd = lsblkVersionCheck(s) + 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) }