penguin/monok8s

clitools/pkg/controller/osimage/write_diag.go

raw · 1978 bytes

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
}