github.com/balzaczyy/golucene@v0.0.0-20151210033525-d0be9ee89713/test_framework/util/throttled.go (about)

     1  package util
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/balzaczyy/golucene/core/store"
     6  	"github.com/balzaczyy/golucene/core/util"
     7  	"time"
     8  )
     9  
    10  // util/ThrottledIndexOutput.java
    11  
    12  const DEFAULT_MIN_WRITTEN_BYTES = 024
    13  
    14  // Intentionally slow IndexOutput for testing.
    15  type ThrottledIndexOutput struct {
    16  	*store.IndexOutputImpl
    17  	bytesPerSecond   int
    18  	delegate         store.IndexOutput
    19  	flushDelayMillis int64
    20  	closeDelayMillis int64
    21  	seekDelayMillis  int64
    22  	pendingBytes     int64
    23  	minBytesWritten  int64
    24  	timeElapsed      int64
    25  	bytes            []byte
    26  }
    27  
    28  func (out *ThrottledIndexOutput) NewFromDelegate(output store.IndexOutput) *ThrottledIndexOutput {
    29  	ans := &ThrottledIndexOutput{
    30  		delegate:         output,
    31  		bytesPerSecond:   out.bytesPerSecond,
    32  		flushDelayMillis: out.flushDelayMillis,
    33  		closeDelayMillis: out.closeDelayMillis,
    34  		seekDelayMillis:  out.seekDelayMillis,
    35  		minBytesWritten:  out.minBytesWritten,
    36  		bytes:            make([]byte, 1),
    37  	}
    38  	ans.IndexOutputImpl = store.NewIndexOutput(ans)
    39  	return ans
    40  }
    41  
    42  func NewThrottledIndexOutput(bytesPerSecond int, delayInMillis int64, delegate store.IndexOutput) *ThrottledIndexOutput {
    43  	assert(bytesPerSecond > 0)
    44  	ans := &ThrottledIndexOutput{
    45  		delegate:         delegate,
    46  		bytesPerSecond:   bytesPerSecond,
    47  		flushDelayMillis: delayInMillis,
    48  		closeDelayMillis: delayInMillis,
    49  		seekDelayMillis:  delayInMillis,
    50  		minBytesWritten:  DEFAULT_MIN_WRITTEN_BYTES,
    51  		bytes:            make([]byte, 1),
    52  	}
    53  	ans.IndexOutputImpl = store.NewIndexOutput(ans)
    54  	return ans
    55  }
    56  
    57  func MBitsToBytes(mBits int) int {
    58  	return mBits * 125000
    59  }
    60  
    61  func (out *ThrottledIndexOutput) Close() error {
    62  	<-time.After(time.Duration(out.closeDelayMillis + out.delay(true)))
    63  	return out.delegate.Close()
    64  }
    65  
    66  func (out *ThrottledIndexOutput) FilePointer() int64 {
    67  	return out.delegate.FilePointer()
    68  }
    69  
    70  func (out *ThrottledIndexOutput) WriteByte(b byte) error {
    71  	out.bytes[0] = b
    72  	return out.WriteBytes(out.bytes)
    73  }
    74  
    75  func (out *ThrottledIndexOutput) WriteBytes(buf []byte) error {
    76  	before := time.Now()
    77  	// TODO: sometimes, write only half the bytes, then sleep, then 2nd
    78  	// half, then sleep, so we sometimes interrupt having only written
    79  	// not all bytes
    80  	err := out.delegate.WriteBytes(buf)
    81  	if err != nil {
    82  		return err
    83  	}
    84  	out.timeElapsed += int64(time.Now().Sub(before))
    85  	out.pendingBytes += int64(len(buf))
    86  	<-time.After(time.Duration(out.delay(false)))
    87  	return nil
    88  }
    89  
    90  func (out *ThrottledIndexOutput) delay(closing bool) int64 {
    91  	if out.pendingBytes > 0 && (closing || out.pendingBytes > out.minBytesWritten) {
    92  		actualBps := (out.timeElapsed / out.pendingBytes) * 1000000000
    93  		if actualBps > int64(out.bytesPerSecond) {
    94  			expected := out.pendingBytes * 1000 / int64(out.bytesPerSecond)
    95  			delay := expected - (out.timeElapsed / 1000000)
    96  			out.pendingBytes = 0
    97  			out.timeElapsed = 0
    98  			return delay
    99  		}
   100  	}
   101  	return 0
   102  }
   103  
   104  func (out *ThrottledIndexOutput) CopyBytes(input util.DataInput, numBytes int64) error {
   105  	return out.delegate.CopyBytes(input, numBytes)
   106  }
   107  
   108  func assert(ok bool) {
   109  	assert2(ok, "assert fail")
   110  }
   111  
   112  func assert2(ok bool, msg string, args ...interface{}) {
   113  	if !ok {
   114  		panic(fmt.Sprintf(msg, args...))
   115  	}
   116  }