package osimage import ( "context" "fmt" "io" "os" "time" ) const ( defaultWriteBufferSize = 1 * 1024 * 1024 defaultMinWriteBPS = int64(2 * 1024 * 1024) defaultInitialWriteBPS = int64(4 * 1024 * 1024) defaultMaxWriteBPS = int64(8 * 1024 * 1024) defaultBurstBytes = int64(512 * 1024) defaultSampleInterval = 250 * time.Millisecond defaultSyncEveryBytes = 0 defaultBusyHighPct = 80.0 defaultBusyLowPct = 40.0 defaultSlowAwait = 20 * time.Millisecond defaultFastAwait = 5 * time.Millisecond ) func WriteStreamToTarget(ctx context.Context, src io.Reader, targetPath string, expectedSize int64, bufferSize int, progress ProgressFunc, ) (int64, error) { if targetPath == "" { return 0, fmt.Errorf("target path is required") } if bufferSize <= 0 { bufferSize = defaultWriteBufferSize } f, err := os.OpenFile(targetPath, os.O_WRONLY, 0) if err != nil { return 0, fmt.Errorf("open target: %w", err) } defer f.Close() ctrl, err := newAdaptiveWriteController(targetPath) if err != nil { ctrl = newNoopAdaptiveWriteController() } written, err := copyWithProgressBuffer( ctx, f, src, expectedSize, "flash", progress, make([]byte, bufferSize), ctrl, defaultSyncEveryBytes, ) if err != nil { return written, err } if expectedSize > 0 && written != expectedSize { return written, fmt.Errorf("written size mismatch: got %d want %d", written, expectedSize) } if err := f.Sync(); err != nil { return written, fmt.Errorf("sync target: %w", err) } return written, nil } func copyWithProgressBuffer( ctx context.Context, dst *os.File, src io.Reader, total int64, stage string, progress ProgressFunc, buf []byte, ctrl *adaptiveWriteController, syncEvery int64, ) (int64, error) { var written int64 var sinceSync int64 for { select { case <-ctx.Done(): return written, ctx.Err() default: } nr, er := src.Read(buf) if nr > 0 { if ctrl != nil { if err := ctrl.Wait(ctx, nr); err != nil { return written, err } } nw, ew := dst.Write(buf[:nr]) if nw > 0 { written += int64(nw) sinceSync += int64(nw) if ctrl != nil { ctrl.ObserveWrite(nw) } if progress != nil { progress(Progress{ Stage: stage, BytesComplete: written, BytesTotal: total, }) } if syncEvery > 0 && sinceSync >= syncEvery { if err := dst.Sync(); err != nil { return written, fmt.Errorf("periodic sync target: %w", err) } sinceSync = 0 if ctrl != nil { ctrl.ObserveSync() } } } if ew != nil { return written, ew } if nw != nr { return written, io.ErrShortWrite } } if er != nil { if er == io.EOF { return written, nil } return written, fmt.Errorf("copy %s: %w", stage, er) } } }