github.com/Cloud-Foundations/Dominator@v0.3.4/lib/rateio/reader.go (about)

     1  package rateio
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"time"
     7  
     8  	"github.com/Cloud-Foundations/Dominator/lib/format"
     9  )
    10  
    11  const (
    12  	DEFAULT_SPEED_PERCENT = 2
    13  )
    14  
    15  func newReaderContext(maxIOPerSecond uint64, speedPercent uint64,
    16  	measurer ReadIOMeasurer) *ReaderContext {
    17  	var ctx ReaderContext
    18  	ctx.maxIOPerSecond = maxIOPerSecond
    19  	if speedPercent < 1 {
    20  		speedPercent = DEFAULT_SPEED_PERCENT
    21  	}
    22  	ctx.speedPercent = speedPercent
    23  	ctx.chunklen = ctx.maxIOPerSecond * ctx.speedPercent / 10000
    24  	ctx.measurer = measurer
    25  	ctx.timeOfLastPause = time.Now()
    26  	measurer.Reset()
    27  	return &ctx
    28  }
    29  
    30  func (ctx *ReaderContext) initialiseMaximumSpeed(maxSpeed uint64) {
    31  	if ctx.maxIOPerSecond > 0 {
    32  		fmt.Println("Maximum speed already set")
    33  	}
    34  	ctx.maxIOPerSecond = maxSpeed
    35  	ctx.chunklen = ctx.maxIOPerSecond * ctx.speedPercent / 10000
    36  }
    37  
    38  func (ctx *ReaderContext) setSpeedPercent(percent uint) {
    39  	if percent > 100 {
    40  		percent = 100
    41  	}
    42  	ctx.speedPercent = uint64(percent)
    43  	ctx.chunklen = ctx.maxIOPerSecond * ctx.speedPercent / 10000
    44  	ctx.timeOfLastPause = time.Now()
    45  	ctx.measurer.Reset()
    46  }
    47  
    48  func (ctx *ReaderContext) newReader(rd io.Reader) *Reader {
    49  	var reader Reader
    50  	reader.ctx = ctx
    51  	reader.rd = rd
    52  	return &reader
    53  }
    54  
    55  func (ctx *ReaderContext) format() string {
    56  	return fmt.Sprintf("max speed=%s/s limit=%d%% %s/s",
    57  		format.FormatBytes(ctx.maxIOPerSecond),
    58  		ctx.speedPercent,
    59  		format.FormatBytes(ctx.maxIOPerSecond*ctx.speedPercent/100))
    60  }
    61  
    62  func (rd *Reader) read(b []byte) (n int, err error) {
    63  	if rd.ctx.maxIOPerSecond < 1 {
    64  		// Unspecified capacity: go at maximum speed.
    65  		return rd.rd.Read(b)
    66  	}
    67  	if rd.ctx.speedPercent >= 100 {
    68  		// Operate at maximum speed: get out of the way.
    69  		return rd.rd.Read(b)
    70  	}
    71  	if rd.ctx.bytesSinceLastPause >= rd.ctx.chunklen {
    72  		// Need to slow down.
    73  		desiredPerSecond := rd.ctx.maxIOPerSecond * rd.ctx.speedPercent / 100
    74  		if desiredPerSecond < 1 {
    75  			desiredPerSecond = rd.ctx.maxIOPerSecond / 1000
    76  		}
    77  		if desiredPerSecond < 1 {
    78  			desiredPerSecond = 1
    79  		}
    80  		readSinceLastPause, err := rd.ctx.measurer.MeasureReadIO(
    81  			rd.ctx.bytesSinceLastPause)
    82  		if err != nil {
    83  			return 0, err
    84  		}
    85  		desiredDuration := time.Duration(uint64(time.Second) *
    86  			readSinceLastPause / desiredPerSecond)
    87  		targetTime := rd.ctx.timeOfLastPause.Add(desiredDuration)
    88  		rd.ctx.timeOfLastPause = time.Now()
    89  		duration := targetTime.Sub(time.Now())
    90  		if duration > 0 {
    91  			if rd.ctx.sleepTimeDistribution != nil {
    92  				rd.ctx.sleepTimeDistribution.Add(duration)
    93  			}
    94  			time.Sleep(duration)
    95  		}
    96  		rd.ctx.bytesSinceLastPause = 0
    97  	}
    98  	n, err = rd.rd.Read(b)
    99  	if n < 1 || err != nil {
   100  		return
   101  	}
   102  	rd.ctx.bytesSinceLastPause += uint64(n)
   103  	return
   104  }