github.com/coocood/badger@v1.5.1-0.20200528065104-c02ac3616d04/fileutil/writer.go (about)

     1  package fileutil
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  
     7  	"github.com/ncw/directio"
     8  	"golang.org/x/time/rate"
     9  )
    10  
    11  // DirectWriter writes to a file opened with O_DIRECT flag.
    12  // `Finish` must be called when the writing is done to truncate and sync the file.
    13  type DirectWriter struct {
    14  	writer
    15  }
    16  
    17  // BufferedWriter writes to a file with buffer.
    18  type BufferedWriter struct {
    19  	writer
    20  }
    21  
    22  type writer struct {
    23  	fd       *os.File
    24  	fileOff  int64
    25  	writeBuf []byte
    26  	bufOff   int64
    27  	limiter  *rate.Limiter
    28  }
    29  
    30  func NewBufferedWriter(fd *os.File, bufSize int, limiter *rate.Limiter) *BufferedWriter {
    31  	return &BufferedWriter{
    32  		writer: writer{
    33  			fd:       fd,
    34  			writeBuf: make([]byte, bufSize),
    35  			limiter:  limiter,
    36  		},
    37  	}
    38  }
    39  
    40  func NewDirectWriter(fd *os.File, bufSize int, limiter *rate.Limiter) *DirectWriter {
    41  	return &DirectWriter{
    42  		writer: writer{
    43  			fd:       fd,
    44  			writeBuf: directio.AlignedBlock(bufSize),
    45  			limiter:  limiter,
    46  		},
    47  	}
    48  }
    49  
    50  func (l *writer) Reset(fd *os.File) {
    51  	l.fd = fd
    52  	l.fileOff = 0
    53  	l.bufOff = 0
    54  }
    55  
    56  func (l *writer) Write(p []byte) (n int, err error) {
    57  	return len(p), l.Append(p)
    58  }
    59  
    60  func (l *writer) Append(val []byte) error {
    61  	for {
    62  		n := copy(l.writeBuf[l.bufOff:], val)
    63  		l.bufOff += int64(n)
    64  		if n == len(val) {
    65  			return nil
    66  		}
    67  		err := l.flush()
    68  		if err != nil {
    69  			return err
    70  		}
    71  		val = val[n:]
    72  	}
    73  }
    74  
    75  func (l *writer) waitRateLimiter() {
    76  	if l.limiter != nil {
    77  		err := l.limiter.WaitN(context.Background(), int(l.bufOff))
    78  		if err != nil {
    79  			panic(err)
    80  		}
    81  	}
    82  }
    83  
    84  func (l *writer) flush() error {
    85  	if l.bufOff == 0 {
    86  		return nil
    87  	}
    88  	l.waitRateLimiter()
    89  	_, err := l.fd.Write(l.writeBuf[:l.bufOff])
    90  	if err != nil {
    91  		return err
    92  	}
    93  	l.fileOff += l.bufOff
    94  	l.bufOff = 0
    95  	return nil
    96  }
    97  
    98  func (l *writer) Offset() int64 {
    99  	return l.fileOff + l.bufOff
   100  }
   101  
   102  func (l *BufferedWriter) Flush() error {
   103  	return l.flush()
   104  }
   105  
   106  func (l *BufferedWriter) Sync() error {
   107  	return Fdatasync(l.fd)
   108  }
   109  
   110  func (l *DirectWriter) Finish() error {
   111  	if l.bufOff == 0 {
   112  		return nil
   113  	}
   114  	finalLength := l.fileOff + l.bufOff
   115  	l.bufOff = alignedSize(l.bufOff)
   116  	err := l.flush()
   117  	if err != nil {
   118  		return err
   119  	}
   120  	l.fileOff = finalLength
   121  	err = l.fd.Truncate(finalLength)
   122  	if err != nil {
   123  		return err
   124  	}
   125  	return Fdatasync(l.fd)
   126  }
   127  
   128  func alignedSize(n int64) int64 {
   129  	return (n + directio.BlockSize - 1) / directio.BlockSize * directio.BlockSize
   130  }