chaosblade-exec-os项目的burnio.go文件解读

#################################################

代码位置：https://github.com/chaosblade-io/chaosblade-exec-os.git

文件位置：chaosblade-exec-os/exec/bin/burnio/burnio.go

这个文件刚开始读起来有点难以理解，甚至懵逼，为何呢？

比如一个可执行文件叫chaos_burnio，这个可执行文件的内容就是第二次以执行脚本的命令的方式执行相同的可执行文件，一开始不理解为何这样写，这就像套娃一样，以为这会进入到一个死循环

然后，这两次调用并不相同，主要在于参数不同，第一次调用会重新组装参数，而且会根据条件来选择执行，因此两者并不会进入死循环，不过呢，第一次看到还是感觉蛮伤脑的，这样组织代码的还是第一次见

# 这是第二次执行：调用可执行文件chaos_burnio，但是参数不再有--start，--stop了，总共四个分支执行，只有--start和--stop才会有第二次执行相同可执行命令，故不会在进入到死循环    
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", script+" "+args)
    
output, err := cmd.CombinedOutput()

/*
 * Copyright 1999-2020 Alibaba Group Holding Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
    "context"
    "flag"
    "fmt"
    "path"
    "strings"
    "time"

    "github.com/chaosblade-io/chaosblade-spec-go/channel"
    "github.com/chaosblade-io/chaosblade-spec-go/util"

    "github.com/chaosblade-io/chaosblade-exec-os/exec"
    "github.com/chaosblade-io/chaosblade-exec-os/exec/bin"
)

const count = 100

var burnIODirectory, burnIOSize string
var burnIORead, burnIOWrite, burnIOStart, burnIOStop, burnIONohup bool

func main() {
    flag.StringVar(&burnIODirectory, "directory", "", "the directory where the disk is burning")
    flag.StringVar(&burnIOSize, "size", "", "block size")
    flag.BoolVar(&burnIOWrite, "write", false, "write io")
    flag.BoolVar(&burnIORead, "read", false, "read io")
    flag.BoolVar(&burnIOStart, "start", false, "start burn io")
    flag.BoolVar(&burnIOStop, "stop", false, "stop burn io")
    flag.BoolVar(&burnIONohup, "nohup", false, "start by nohup")
    bin.ParseFlagAndInitLog()

    if burnIOStart {
        startBurnIO(burnIODirectory, burnIOSize, burnIORead, burnIOWrite)
    } else if burnIOStop {
        stopBurnIO(burnIODirectory, burnIORead, burnIOWrite)
    } else if burnIONohup {
        if burnIORead {
            go burnRead(burnIODirectory, burnIOSize)
        }
        if burnIOWrite {
            go burnWrite(burnIODirectory, burnIOSize)
        }
        select {}
    } else {
        bin.PrintErrAndExit("less --start or --stop flag")
    }
}

var readFile = "chaos_burnio.read"
var writeFile = "chaos_burnio.write"
var burnIOBin = exec.BurnIOBin
var logFile = util.GetNohupOutput(util.Bin, "chaos_burnio.log")

var cl = channel.NewLocalChannel()

var stopBurnIOFunc = stopBurnIO

// start burn io
func startBurnIO(directory, size string, read, write bool) {
    ctx := context.Background()
    response := cl.Run(ctx, "nohup",
        fmt.Sprintf(`%s --directory %s --size %s --read=%t --write=%t --nohup=true > %s 2>&1 &`,
            path.Join(util.GetProgramPath(), burnIOBin), directory, size, read, write, logFile))
    if !response.Success {
        stopBurnIOFunc(directory, read, write)
        bin.PrintErrAndExit(response.Err)
        return
    }
    // check
    time.Sleep(time.Second)
    response = cl.Run(ctx, "grep", fmt.Sprintf("%s %s", bin.ErrPrefix, logFile))
    if response.Success {
        errMsg := strings.TrimSpace(response.Result.(string))
        if errMsg != "" {
            stopBurnIOFunc(directory, read, write)
            bin.PrintErrAndExit(errMsg)
            return
        }
    }
    bin.PrintOutputAndExit("success")
}

// stop burn io,  no need to add os.Exit
func stopBurnIO(directory string, read, write bool) {
    ctx := context.WithValue(context.Background(), channel.ExcludeProcessKey, "--stop")
    if read {
        // dd process
        pids, _ := cl.GetPidsByProcessName(readFile, ctx)
        if pids != nil && len(pids) > 0 {
            cl.Run(ctx, "kill", fmt.Sprintf("-9 %s", strings.Join(pids, " ")))
        }
        // chaos_burnio process
        ctxWithKey := context.WithValue(ctx, channel.ProcessKey, burnIOBin)
        pids, _ = cl.GetPidsByProcessName("--read=true", ctxWithKey)
        if pids != nil && len(pids) > 0 {
            cl.Run(ctx, "kill", fmt.Sprintf("-9 %s", strings.Join(pids, " ")))
        }
        cl.Run(ctx, "rm", fmt.Sprintf("-rf %s*", path.Join(directory, readFile)))
    }
    if write {
        // dd process
        pids, _ := cl.GetPidsByProcessName(writeFile, ctx)
        if pids != nil && len(pids) > 0 {
            cl.Run(ctx, "kill", fmt.Sprintf("-9 %s", strings.Join(pids, " ")))
        }
        ctxWithKey := context.WithValue(ctx, channel.ProcessKey, burnIOBin)
        pids, _ = cl.GetPidsByProcessName("--write=true", ctxWithKey)
        if pids != nil && len(pids) > 0 {
            cl.Run(ctx, "kill", fmt.Sprintf("-9 %s", strings.Join(pids, " ")))
        }
        cl.Run(ctx, "rm", fmt.Sprintf("-rf %s*", path.Join(directory, writeFile)))
    }
}

// write burn
func burnWrite(directory, size string) {
    tmpFileForWrite := path.Join(directory, writeFile)
    for {
        args := fmt.Sprintf(`if=/dev/zero of=%s bs=%sM count=%d oflag=dsync`, tmpFileForWrite, size, count)
        response := cl.Run(context.TODO(), "dd", args)
        if !response.Success {
            bin.PrintAndExitWithErrPrefix(response.Err)
            return
        }
    }
}

// read burn
func burnRead(directory, size string) {
    // create a 600M file under the directory
    tmpFileForRead := path.Join(directory, readFile)
    createArgs := fmt.Sprintf("if=/dev/zero of=%s bs=%dM count=%d oflag=dsync", tmpFileForRead, 6, count)
    response := cl.Run(context.TODO(), "dd", createArgs)
    if !response.Success {
        bin.PrintAndExitWithErrPrefix(
            fmt.Sprintf("using dd command to create a temp file under %s directory for reading error, %s",
                directory, response.Err))
    }
    for {
        args := fmt.Sprintf(`if=%s of=/dev/null bs=%sM count=%d iflag=dsync,direct,fullblock`, tmpFileForRead, size, count)
        response = cl.Run(context.TODO(), "dd", args)
        if !response.Success {
            bin.PrintAndExitWithErrPrefix(fmt.Sprintf("using dd command to burn read io error, %s", response.Err))
            return
        }
    }
}

##################################################