Fix bug and implement recover for disk chaos (#73)

* fix bugs in SplitBytesByProcessNum & add overwrite control

Signed-off-by: Andrewmatilde <davis6813585853062@outlook.com>

* delete overwrite control

Signed-off-by: Andrewmatilde <davis6813585853062@outlook.com>

* add recover for disk&&delete fill destory

Signed-off-by: Andrewmatilde <davis6813585853062@outlook.com>

* add & fix comments

Signed-off-by: Andrewmatilde <davis6813585853062@outlook.com>

* prevent user from writing in an existing file.

Signed-off-by: Andrewmatilde <davis6813585853062@outlook.com>

cmd/attack/disk.go

@@ -126,7 +126,6 @@ func NewDiskFillCommand(dep fx.Option, options *core.DiskOption) *cobra.Command
 	cmd.Flags().StringVarP(&options.Percent, "percent", "c", "",
 		"'percent' how many percent data of disk will fill in the file path")
 	cmd.Flags().BoolVarP(&options.FillByFallocate, "fallocate", "f", true, "fill disk by fallocate instead of dd")
-	cmd.Flags().BoolVarP(&options.DestroyFile, "destroy", "d", false, "destroy file after filled in or allocated")
 	return cmd
 }
 

cmd/attack/disk_test.go

@@ -47,7 +47,7 @@ func TestServer_DiskFill(t *testing.T) {
 							Kind:   core.DiskAttack,
 						},
 						Size:              "1024M",
-						Path:              "temp",
+						Path:              "./temp",
 						FillByFallocate:   true,
 						PayloadProcessNum: 1,
 					},
@@ -60,7 +60,7 @@ func TestServer_DiskFill(t *testing.T) {
 							Kind:   core.DiskAttack,
 						},
 						Size:              "24MB",
-						Path:              "temp",
+						Path:              "./temp",
 						FillByFallocate:   false,
 						PayloadProcessNum: 1,
 					},
@@ -71,15 +71,7 @@ func TestServer_DiskFill(t *testing.T) {
 		fx.Invoke(func(s *chaosd.Server, tests []diskTest) {
 			for _, tt := range tests {
 				t.Run(tt.name, func(t *testing.T) {
-					f, err := os.Create(tt.option.Path)
+					_, err := s.ExecuteAttack(chaosd.DiskAttack, tt.option, core.CommandMode)
-					if err != nil {
-						t.Errorf("unexpected err %v when creating temp file", err)
-						return
-					}
-					if f != nil {
-						_ = f.Close()
-					}
-					_, err = s.ExecuteAttack(chaosd.DiskAttack, tt.option, core.CommandMode)
 					if (err != nil) != tt.wantErr {
 						t.Errorf("DiskFill() error = %v, wantErr %v", err, tt.wantErr)
 						return
@@ -116,7 +108,7 @@ func TestServer_DiskPayload(t *testing.T) {
 							Kind:   core.DiskAttack,
 						},
 						Size:              "24M",
-						Path:              "temp",
+						Path:              "./temp",
 						PayloadProcessNum: 1,
 					},
 					wantErr: false,
@@ -128,7 +120,7 @@ func TestServer_DiskPayload(t *testing.T) {
 							Kind:   core.DiskAttack,
 						},
 						Size:              "24M",
-						Path:              "temp",
+						Path:              "./temp",
 						PayloadProcessNum: 1,
 					},
 					wantErr: false,
@@ -138,24 +130,20 @@ func TestServer_DiskPayload(t *testing.T) {
 		fx.Invoke(func(s *chaosd.Server, tests []diskTest) {
 			for _, tt := range tests {
 				t.Run(tt.name, func(t *testing.T) {
-					f, err := os.Create(tt.option.Path)
+					_, err := s.ExecuteAttack(chaosd.DiskAttack, &core.DiskOption{
-					if err != nil {
-						t.Errorf("unexpected err %v when creating temp file", err)
-						return
-					}
-					if f != nil {
-						_ = f.Close()
-					}
-
-					_, err = s.ExecuteAttack(chaosd.DiskAttack, &core.DiskOption{
 						CommonAttackConfig: core.CommonAttackConfig{
 							Action: core.DiskFillAction,
 							Kind:   core.DiskAttack,
 						},
-						Size:            tt.option.Size,
+						PayloadProcessNum: 1,
-						Path:            "temp",
+						Size:              tt.option.Size,
-						FillByFallocate: true,
+						Path:              "./temp",
+						FillByFallocate:   true,
 					}, core.CommandMode)
+					if err != nil {
+						t.Error(err)
+						return
+					}
 					_, err = s.ExecuteAttack(chaosd.DiskAttack, tt.option, core.CommandMode)
 					if (err != nil) != tt.wantErr {
 						t.Errorf("DiskPayload() error = %v, wantErr %v", err, tt.wantErr)
@@ -185,7 +173,6 @@ func writeArgsToDiskOption(args writeArgs) core.DiskOption {
 		Path:              args.Path,
 		Percent:           "",
 		FillByFallocate:   false,
-		DestroyFile:       false,
 		PayloadProcessNum: args.PayloadProcessNum,
 	}
 }
@@ -272,7 +259,6 @@ func readArgsToDiskOption(args readArgs) core.DiskOption {
 		Path:              args.Path,
 		Percent:           "",
 		FillByFallocate:   false,
-		DestroyFile:       false,
 		PayloadProcessNum: args.PayloadProcessNum,
 	}
 }

pkg/core/disk.go

@@ -16,6 +16,8 @@ package core
 import (
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
+	"os"
 	"strconv"
 
 	"github.com/chaos-mesh/chaosd/pkg/utils"
@@ -33,9 +35,9 @@ type DiskOption struct {
 	Size              string `json:"size"`
 	Path              string `json:"path"`
 	Percent           string `json:"percent"`
-	FillByFallocate   bool   `json:"fill_by_fallocate"`
-	DestroyFile       bool   `json:"destroy_file"`
 	PayloadProcessNum uint8  `json:"payload_process_num"`
+
+	FillByFallocate bool `json:"fill_by_fallocate"`
 }
 
 var _ AttackConfig = &DiskOption{}
@@ -58,11 +60,35 @@ func (d *DiskOption) Validate() error {
 			return fmt.Errorf("unknown units of size : %s, DiskOption : %v", d.Size, d)
 		}
 	}
-	if d.Action == DiskFillAction {
+
-		if d.FillByFallocate && byteSize == 0 {
+	if d.Action == DiskFillAction || d.Action == DiskWritePayloadAction {
+		if d.Action == DiskFillAction && d.FillByFallocate && byteSize == 0 {
 			return fmt.Errorf("fallocate not suppurt 0 size or 0 percent data, "+
 				"if you want allocate a 0 size file please set fallocate=false, DiskOption : %v", d)
 		}
+
+		_, err := os.Stat(d.Path)
+		if err != nil {
+			if os.IsNotExist(err) {
+				// check if Path of file is valid when Path is not empty
+				if d.Path != "" {
+					var b []byte
+					if err := ioutil.WriteFile(d.Path, b, 0644); err != nil {
+						return err
+					}
+					if err := os.Remove(d.Path); err != nil {
+						return err
+					}
+				}
+			} else {
+				return err
+			}
+		} else {
+			if d.Action == DiskFillAction {
+				return fmt.Errorf("fill into an existing file")
+			}
+			return fmt.Errorf("write into an existing file")
+		}
 	}
 
 	if d.PayloadProcessNum == 0 {

pkg/server/chaosd/disk.go

@@ -84,12 +84,6 @@ func (diskAttack) diskPayload(payload *core.DiskOption) error {
 			if err != nil {
 				return err
 			}
-			defer func() {
-				err := os.Remove(payload.Path)
-				if err != nil {
-					log.Error(fmt.Sprintf("unexpected err when removing temp file %s", payload.Path), zap.Error(err))
-				}
-			}()
 		}
 	case core.DiskReadPayloadAction:
 		cmdFormat = DDReadPayloadCommand
@@ -157,7 +151,8 @@ func (diskAttack) diskPayload(payload *core.DiskOption) error {
 	return nil
 }
 
-const DDFillCommand = "dd if=/dev/zero of=%s bs=%s count=%s iflag=fullblock"
+// dd command with 'oflag=append conv=notrunc' will append new data in the file.
+const DDFillCommand = "dd if=/dev/zero of=%s bs=%s count=%s iflag=fullblock oflag=append conv=notrunc"
 const FallocateCommand = "fallocate -l %s %s"
 
 // diskFill will execute a dd command (DDFillCommand or FallocateCommand)
@@ -172,15 +167,6 @@ func (diskAttack) diskFill(fill *core.DiskOption) error {
 		}
 	}
 
-	defer func() {
-		if fill.DestroyFile {
-			err := os.Remove(fill.Path)
-			if err != nil {
-				log.Error(fmt.Sprintf("unexpected err when removing file %s", fill.Path), zap.Error(err))
-			}
-		}
-	}()
-
 	if fill.Size != "" {
 		fill.Size = strings.Trim(fill.Size, " ")
 	} else if fill.Percent != "" {
@@ -196,29 +182,56 @@ func (diskAttack) diskFill(fill *core.DiskOption) error {
 			log.Error("fail to get disk total size", zap.Error(err))
 			return err
 		}
-		fill.Size = strconv.FormatUint(totalSize*percent/100, 10)
+		fill.Size = strconv.FormatUint(totalSize*percent/100, 10) + "c"
 	}
-
 	var cmd *exec.Cmd
 	if fill.FillByFallocate {
 		cmd = exec.Command("bash", "-c", fmt.Sprintf(FallocateCommand, fill.Size, fill.Path))
-	} else {
+		output, err := cmd.CombinedOutput()
-		//1Unit means the block size. The bytes size dd read | write is (block size) * (size).
+		if err != nil {
-		cmd = exec.Command("bash", "-c", fmt.Sprintf(DDFillCommand, fill.Path, fill.Size, "1"))
+			log.Error(string(output), zap.Error(err))
-	}
+			return err
-
+		}
-	output, err := cmd.CombinedOutput()
-
-	if err != nil {
-		log.Error(string(output), zap.Error(err))
-	} else {
 		log.Info(string(output))
+	} else {
+		byteSize, err := utils.ParseUnit(fill.Size)
+		if err != nil {
+			log.Error("fail to parse disk size", zap.Error(err))
+			return err
+		}
+
+		ddBlocks, err := utils.SplitBytesByProcessNum(byteSize, 1)
+		if err != nil {
+			log.Error("fail to split disk size", zap.Error(err))
+			return err
+		}
+		for _, block := range ddBlocks {
+			cmd = exec.Command("bash", "-c", fmt.Sprintf(DDFillCommand, fill.Path, block.BlockSize, block.Count))
+			output, err := cmd.CombinedOutput()
+
+			if err != nil {
+				log.Error(string(output), zap.Error(err))
+				return err
+			}
+			log.Info(string(output))
+		}
 	}
 
-	return err
+	return nil
 }
 
 func (diskAttack) Recover(exp core.Experiment, _ Environment) error {
-	log.Info("Recover disk attack will do nothing, because delete | truncate data is too dangerous.")
+	config, err := exp.GetRequestCommand()
+	if err != nil {
+		return err
+	}
+	option := *config.(*core.DiskOption)
+	switch option.Action {
+	case core.DiskFillAction, core.DiskWritePayloadAction:
+		err = os.Remove(option.Path)
+		if err != nil {
+			log.Warn(fmt.Sprintf("recover disk: remove %s failed", option.Path), zap.Error(err))
+		}
+	}
 	return nil
 }

pkg/utils/units.go

@@ -88,9 +88,17 @@ func SplitBytesByProcessNum(bytes uint64, processNum uint8) ([]DdArgBlock, error
 			bytes -= blockSize
 		}
 	}
-	ddArgBlocks = append(ddArgBlocks, DdArgBlock{
+
-		BlockSize: "1",
+	if bytes == 0 {
-		Count:     strconv.FormatUint(bytes, 10) + "c",
+		ddArgBlocks = append(ddArgBlocks, DdArgBlock{
-	})
+			Count:     "0",
+			BlockSize: "1M",
+		})
+	} else {
+		ddArgBlocks = append(ddArgBlocks, DdArgBlock{
+			Count:     "1",
+			BlockSize: strconv.FormatUint(bytes, 10) + "c",
+		})
+	}
 	return ddArgBlocks, nil
 }