failureinjection: add disk stall cycle option

Some usages of disk stalls wish to repeatedly stall
and unstall a disk in a short period of time. The current
implementation requires an ssh connection per stall/unstall
call, which can easily lead to flakes. e.g. the unstall
call runs into an ssh flake and the node fatals before
it can be recovered.

This change adds a disk stall cycle option which will
do this in one single ssh connection until Recover is
called.

Author: DarrylWong

pkg/cmd/roachtest/roachtestutil/disk_stall.go

@@ -7,6 +7,7 @@ package roachtestutil
 
 import (
 	"context"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
@@ -22,6 +23,7 @@ type DiskStaller interface {
 	Setup(ctx context.Context)
 	Cleanup(ctx context.Context)
 	Stall(ctx context.Context, nodes option.NodeListOption)
+	StallCycle(ctx context.Context, nodes option.NodeListOption, stallDuration, unstallDuration time.Duration)
 	Slow(ctx context.Context, nodes option.NodeListOption, bytesPerSecond int)
 	Unstall(ctx context.Context, nodes option.NodeListOption)
 	DataDir() string
@@ -32,13 +34,16 @@ type NoopDiskStaller struct{}
 
 var _ DiskStaller = NoopDiskStaller{}
 
-func (n NoopDiskStaller) Cleanup(ctx context.Context)                            {}
-func (n NoopDiskStaller) DataDir() string                                        { return "{store-dir}" }
-func (n NoopDiskStaller) LogDir() string                                         { return "logs" }
-func (n NoopDiskStaller) Setup(ctx context.Context)                              {}
-func (n NoopDiskStaller) Slow(_ context.Context, _ option.NodeListOption, _ int) {}
-func (n NoopDiskStaller) Stall(_ context.Context, _ option.NodeListOption)       {}
-func (n NoopDiskStaller) Unstall(_ context.Context, _ option.NodeListOption)     {}
+func (n NoopDiskStaller) Cleanup(ctx context.Context)                                               {}
+func (n NoopDiskStaller) DataDir() string                                                           { return "{store-dir}" }
+func (n NoopDiskStaller) LogDir() string                                                            { return "logs" }
+func (n NoopDiskStaller) Setup(ctx context.Context)                                                 {}
+func (n NoopDiskStaller) Slow(_ context.Context, _ option.NodeListOption, _ int)                    {}
+func (n NoopDiskStaller) Stall(_ context.Context, _ option.NodeListOption)                          {}
+func (n NoopDiskStaller) StallCycle(
+	_ context.Context, _ option.NodeListOption, _, _ time.Duration,
+) {}
+func (n NoopDiskStaller) Unstall(_ context.Context, _ option.NodeListOption)                        {}
 
 type Fataler interface {
 	Fatal(args ...interface{})
@@ -103,6 +108,23 @@ func (s *cgroupDiskStaller) Stall(ctx context.Context, nodes option.NodeListOpti
 	}
 }
 
+func (s *cgroupDiskStaller) StallCycle(
+	ctx context.Context, nodes option.NodeListOption, stallDuration, unstallDuration time.Duration,
+) {
+	l := newDiskStallLogger(s.f.L(), nodes, "Stall")
+	if err := s.Failer.Inject(ctx, l, failures.DiskStallArgs{
+		StallLogs:            s.stallLogs,
+		StallWrites:          true,
+		StallReads:           s.stallReads,
+		Nodes:                nodes.InstallNodes(),
+		Cycle:                true,
+		CycleStallDuration:   stallDuration,
+		CycleUnstallDuration: unstallDuration,
+	}); err != nil {
+		s.f.Fatalf("failed to stall disk: %s", err)
+	}
+}
+
 func (s *cgroupDiskStaller) Slow(
 	ctx context.Context, nodes option.NodeListOption, bytesPerSecond int,
 ) {
@@ -168,6 +190,20 @@ func (s *dmsetupDiskStaller) Stall(ctx context.Context, nodes option.NodeListOpt
 	}
 }
 
+func (s *dmsetupDiskStaller) StallCycle(
+	ctx context.Context, nodes option.NodeListOption, stallDuration, unstallDuration time.Duration,
+) {
+	l := newDiskStallLogger(s.f.L(), nodes, "Stall")
+	if err := s.Failer.Inject(ctx, l, failures.DiskStallArgs{
+		Nodes:                nodes.InstallNodes(),
+		Cycle:                true,
+		CycleStallDuration:   stallDuration,
+		CycleUnstallDuration: unstallDuration,
+	}); err != nil {
+		s.f.Fatalf("failed to stall disk: %s", err)
+	}
+}
+
 func (s *dmsetupDiskStaller) Slow(
 	ctx context.Context, nodes option.NodeListOption, bytesPerSecond int,
 ) {

pkg/cmd/roachtest/tests/disk_stall.go

@@ -462,7 +462,11 @@ func runDiskStalledWALFailoverWithProgress(ctx context.Context, t test.Test, c c
 	// Use CgroupDiskStaller with readsToo=false to only stall writes.
 	s := roachtestutil.MakeCgroupDiskStaller(t, c, false /* readsToo */, false /* logsToo */)
 	s.Setup(ctx)
-	defer s.Cleanup(ctx)
+	// NB: We use a background context in the defer'ed cleanup command,
+	// otherwise on test failure our c.Run calls will be ignored. Leaving
+	// the disk stalled will prevent artifact collection, making debugging
+	// difficult.
+	defer s.Cleanup(context.Background())
 
 	t.Status("starting cluster")
 	startOpts := option.DefaultStartOpts()
@@ -612,36 +616,15 @@ func runDiskStalledWALFailoverWithProgress(ctx context.Context, t test.Test, c c
 				case <-time.After(diskStallWaitDur):
 					t.Status("starting disk stall")
 				}
-				stallStart := timeutil.Now()
-				// Execute short 200ms stalls every 5s.
-				for timeutil.Since(stallStart) < operationDur {
-					select {
-					case <-ctx.Done():
-						t.Fatalf("context done while stall induced: %s", ctx.Err())
-					case <-time.After(stallInterval):
-						func() {
-							t.Status("short disk stall on n1")
-							s.Stall(ctx, c.Node(1))
-							defer func() {
-								// NB: We use a background context in the defer'ed unstall command,
-								// otherwise on test failure our Unstall calls will be ignored. Leaving
-								// the disk stalled will prevent artifact collection, making debugging
-								// difficult.
-								ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
-								defer cancel()
-								s.Unstall(ctx, c.Node(1))
-								t.Status("unstalled disk on n1")
-							}()
-							select {
-							case <-ctx.Done():
-								t.Fatalf("context done while stall induced: %s", ctx.Err())
-							case <-time.After(shortStallDur):
-								return
-							}
-						}()
-					}
+				// Execute short 200ms stalls every 5s for 3 minutes.
+				s.StallCycle(ctx, c.Node(1), shortStallDur, stallInterval)
+				select {
+				case <-ctx.Done():
+					t.Fatalf("context done while stall induced: %s", ctx.Err())
+				case <-time.After(operationDur):
+					s.Unstall(ctx, c.Node(1))
+					t.Status("disk stalls stopped")
 				}
-
 				return nil
 			}, task.Name("disk-stall-phase"))
 		} else {