Throttle disk write to prevent etcd puking during upgrade

clitools/pkg/controller/osimage/write_diag.go

@@ -0,0 +1,106 @@
+package osimage
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"time"
+
+	"k8s.io/klog/v2"
+)
+
+type repeatPatternReader struct {
+	pattern []byte
+	remain  int64
+	off     int
+}
+
+func newRepeatPatternReader(total int64, pattern []byte) *repeatPatternReader {
+	if len(pattern) == 0 {
+		pattern = []byte("monok8s-test-pattern-0123456789abcdef")
+	}
+	return &repeatPatternReader{
+		pattern: pattern,
+		remain:  total,
+	}
+}
+
+func (r *repeatPatternReader) Read(p []byte) (int, error) {
+	if r.remain <= 0 {
+		return 0, io.EOF
+	}
+
+	if int64(len(p)) > r.remain {
+		p = p[:r.remain]
+	}
+
+	n := 0
+	for n < len(p) {
+		copied := copy(p[n:], r.pattern[r.off:])
+		n += copied
+		r.off += copied
+		if r.off == len(r.pattern) {
+			r.off = 0
+		}
+	}
+
+	r.remain -= int64(n)
+	return n, nil
+}
+
+func TestStreamToTarget(ctx context.Context, targetPath string) error {
+	const (
+		totalSize  = int64(512 * 1024 * 1024) // 512 MiB
+		bufferSize = 128 * 1024               // test the conservative setting
+	)
+
+	src := newRepeatPatternReader(totalSize, nil)
+
+	start := time.Now()
+	lastLog := start
+
+	progress := func(p Progress) {
+		now := time.Now()
+		if now.Sub(lastLog) < 1*time.Second && p.BytesComplete != p.BytesTotal {
+			return
+		}
+		lastLog = now
+
+		var mbps float64
+		elapsed := now.Sub(start).Seconds()
+		if elapsed > 0 {
+			mbps = float64(p.BytesComplete) / 1024.0 / 1024.0 / elapsed
+		}
+
+		klog.InfoS("test write progress",
+			"stage", p.Stage,
+			"bytesComplete", p.BytesComplete,
+			"bytesTotal", p.BytesTotal,
+			"mbpsAvg", fmt.Sprintf("%.2f", mbps),
+		)
+	}
+
+	written, err := WriteStreamToTarget(
+		ctx,
+		src,
+		targetPath,
+		totalSize,
+		bufferSize,
+		progress,
+	)
+	if err != nil {
+		return fmt.Errorf("write stream to target: %w", err)
+	}
+
+	elapsed := time.Since(start)
+	mbps := float64(written) / 1024.0 / 1024.0 / elapsed.Seconds()
+
+	klog.InfoS("test write complete",
+		"targetPath", targetPath,
+		"written", written,
+		"elapsed", elapsed.String(),
+		"mbpsAvg", fmt.Sprintf("%.2f", mbps),
+	)
+
+	return nil
+}