github.com/etecs-ru/ristretto@v0.9.1/z/buffer.go (about)

     1  /*
     2   * Copyright 2020 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package z
    18  
    19  import (
    20  	"encoding/binary"
    21  	"errors"
    22  	"fmt"
    23  	"log"
    24  	"os"
    25  	"sort"
    26  	"sync/atomic"
    27  )
    28  
    29  const (
    30  	defaultCapacity = 64
    31  	defaultTag      = "buffer"
    32  )
    33  
    34  // Buffer is equivalent of bytes.Buffer without the ability to read. It is NOT thread-safe.
    35  //
    36  // In UseCalloc mode, z.Calloc is used to allocate memory, which depending upon how the code is
    37  // compiled could use jemalloc for allocations.
    38  //
    39  // In UseMmap mode, Buffer  uses file mmap to allocate memory. This allows us to store big data
    40  // structures without using physical memory.
    41  //
    42  // MaxSize can be set to limit the memory usage.
    43  type Buffer struct {
    44  	mmapFile      *MmapFile
    45  	autoMmapDir   string
    46  	tag           string
    47  	buf           []byte
    48  	offset        uint64
    49  	bufType       BufferType
    50  	maxSz         int
    51  	padding       uint64
    52  	autoMmapAfter int
    53  	curSz         int
    54  	persistent    bool
    55  }
    56  
    57  func NewBuffer(capacity int, tag string) *Buffer {
    58  	if capacity < defaultCapacity {
    59  		capacity = defaultCapacity
    60  	}
    61  	if tag == "" {
    62  		tag = defaultTag
    63  	}
    64  	return &Buffer{
    65  		buf:     Calloc(capacity, tag),
    66  		bufType: UseCalloc,
    67  		curSz:   capacity,
    68  		offset:  8,
    69  		padding: 8,
    70  		tag:     tag,
    71  	}
    72  }
    73  
    74  // It is the caller's responsibility to set offset after this, because Buffer
    75  // doesn't remember what it was.
    76  func NewBufferPersistent(path string, capacity int) (*Buffer, error) {
    77  	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, filePerm) //nolint:gosec,revive //adopt fork, do not touch it
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  	buffer, err := newBufferFile(file, capacity)
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  	buffer.persistent = true
    86  	return buffer, nil
    87  }
    88  
    89  func NewBufferTmp(dir string, capacity int) (*Buffer, error) {
    90  	if dir == "" {
    91  		dir = tmpDir
    92  	}
    93  
    94  	file, err := os.CreateTemp(dir, "buffer")
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  	return newBufferFile(file, capacity)
    99  }
   100  
   101  func newBufferFile(file *os.File, capacity int) (*Buffer, error) {
   102  	if capacity < defaultCapacity {
   103  		capacity = defaultCapacity
   104  	}
   105  	mmapFile, err := OpenMmapFileUsing(file, capacity, true)
   106  
   107  	if err != nil && !errors.Is(err, ErrNewFileCreateFailed) {
   108  		return nil, err
   109  	}
   110  
   111  	buf := &Buffer{
   112  		buf:      mmapFile.Data,
   113  		bufType:  UseMmap,
   114  		curSz:    len(mmapFile.Data),
   115  		mmapFile: mmapFile,
   116  		offset:   8,
   117  		padding:  8,
   118  	}
   119  	return buf, nil
   120  }
   121  
   122  func NewBufferSlice(slice []byte) *Buffer {
   123  	return &Buffer{
   124  		offset:  uint64(len(slice)),
   125  		buf:     slice,
   126  		bufType: UseInvalid,
   127  	}
   128  }
   129  
   130  func (b *Buffer) WithAutoMmap(threshold int, path string) *Buffer {
   131  	if b.bufType != UseCalloc {
   132  		panic("can only autoMmap with UseCalloc")
   133  	}
   134  	b.autoMmapAfter = threshold
   135  	if path == "" {
   136  		b.autoMmapDir = tmpDir
   137  	} else {
   138  		b.autoMmapDir = path
   139  	}
   140  	return b
   141  }
   142  
   143  func (b *Buffer) WithMaxSize(size int) *Buffer {
   144  	b.maxSz = size
   145  	return b
   146  }
   147  
   148  func (b *Buffer) IsEmpty() bool {
   149  	return int(b.offset) == b.StartOffset()
   150  }
   151  
   152  // LenWithPadding would return the number of bytes written to the buffer so far
   153  // plus the padding at the start of the buffer.
   154  func (b *Buffer) LenWithPadding() int {
   155  	return int(atomic.LoadUint64(&b.offset))
   156  }
   157  
   158  // LenNoPadding would return the number of bytes written to the buffer so far
   159  // (without the padding).
   160  func (b *Buffer) LenNoPadding() int {
   161  	return int(atomic.LoadUint64(&b.offset) - b.padding)
   162  }
   163  
   164  // Bytes would return all the written bytes as a slice.
   165  func (b *Buffer) Bytes() []byte {
   166  	off := atomic.LoadUint64(&b.offset)
   167  	return b.buf[b.padding:off]
   168  }
   169  
   170  // Grow would grow the buffer to have at least n more bytes. In case the buffer is at capacity, it
   171  // would reallocate twice the size of current capacity + n, to ensure n bytes can be written to the
   172  // buffer without further allocation. In UseMmap mode, this might result in underlying file
   173  // expansion.
   174  func (b *Buffer) Grow(n int) {
   175  	if b.buf == nil {
   176  		panic("z.Buffer needs to be initialized before using")
   177  	}
   178  	if b.maxSz > 0 && int(b.offset)+n > b.maxSz {
   179  		err := fmt.Errorf(
   180  			"z.Buffer max size exceeded: %d offset: %d grow: %d", b.maxSz, b.offset, n)
   181  		panic(err)
   182  	}
   183  	if int(b.offset)+n < b.curSz {
   184  		return
   185  	}
   186  
   187  	// Calculate new capacity.
   188  	growBy := b.curSz + n
   189  	// Don't allocate more than 1GB at a time.
   190  	if growBy > 1<<30 {
   191  		growBy = 1 << 30
   192  	}
   193  	// Allocate at least n, even if it exceeds the 1GB limit above.
   194  	if n > growBy {
   195  		growBy = n
   196  	}
   197  	b.curSz += growBy
   198  
   199  	switch b.bufType {
   200  	case UseCalloc:
   201  		// If autoMmap gets triggered, copy the slice over to an mmaped file.
   202  		if b.autoMmapAfter > 0 && b.curSz > b.autoMmapAfter {
   203  			b.bufType = UseMmap
   204  			file, err := os.CreateTemp(b.autoMmapDir, "")
   205  			if err != nil {
   206  				panic(err)
   207  			}
   208  			mmapFile, err := OpenMmapFileUsing(file, b.curSz, true)
   209  
   210  			if err != nil && !errors.Is(err, ErrNewFileCreateFailed) {
   211  				panic(err)
   212  			}
   213  			assert(int(b.offset) == copy(mmapFile.Data, b.buf[:b.offset]))
   214  			Free(b.buf)
   215  			b.mmapFile = mmapFile
   216  			b.buf = mmapFile.Data
   217  			break
   218  		}
   219  
   220  		// Else, reallocate the slice.
   221  		newBuf := Calloc(b.curSz, b.tag)
   222  		assert(int(b.offset) == copy(newBuf, b.buf[:b.offset]))
   223  		Free(b.buf)
   224  		b.buf = newBuf
   225  
   226  	case UseMmap:
   227  		// Truncate and remap the underlying file.
   228  		if err := b.mmapFile.Truncate(int64(b.curSz)); err != nil {
   229  			panic(fmt.Sprintf("%v while trying to truncate file: %s to size: %d",
   230  				err, b.mmapFile.Fd.Name(), b.curSz))
   231  		}
   232  		b.buf = b.mmapFile.Data
   233  
   234  	default:
   235  		panic("can only use Grow on UseCalloc and UseMmap buffers")
   236  	}
   237  }
   238  
   239  // Allocate is a way to get a slice of size n back from the buffer. This slice can be directly
   240  // written to. Warning: Allocate is not thread-safe. The byte slice returned MUST be used before
   241  // further calls to Buffer.
   242  func (b *Buffer) Allocate(n int) []byte {
   243  	b.Grow(n)
   244  	off := b.offset
   245  	b.offset += uint64(n)
   246  	return b.buf[off:int(b.offset)]
   247  }
   248  
   249  // AllocateOffset works the same way as allocate, but instead of returning a byte slice, it returns
   250  // the offset of the allocation.
   251  func (b *Buffer) AllocateOffset(n int) int {
   252  	b.Grow(n)
   253  	b.offset += uint64(n)
   254  	return int(b.offset) - n
   255  }
   256  
   257  func (b *Buffer) writeLen(sz int) {
   258  	buf := b.Allocate(4)
   259  	binary.BigEndian.PutUint32(buf, uint32(sz))
   260  }
   261  
   262  // SliceAllocate would encode the size provided into the buffer, followed by a call to Allocate,
   263  // hence returning the slice of size sz. This can be used to allocate a lot of small buffers into
   264  // this big buffer.
   265  // Note that SliceAllocate should NOT be mixed with normal calls to Write.
   266  func (b *Buffer) SliceAllocate(sz int) []byte {
   267  	b.Grow(4 + sz)
   268  	b.writeLen(sz)
   269  	return b.Allocate(sz)
   270  }
   271  
   272  func (b *Buffer) StartOffset() int {
   273  	return int(b.padding)
   274  }
   275  
   276  func (b *Buffer) WriteSlice(slice []byte) {
   277  	dst := b.SliceAllocate(len(slice))
   278  	assert(len(slice) == copy(dst, slice))
   279  }
   280  
   281  func (b *Buffer) SliceIterate(f func(slice []byte) error) error {
   282  	if b.IsEmpty() {
   283  		return nil
   284  	}
   285  	slice, next := []byte{}, b.StartOffset()
   286  	for next >= 0 {
   287  		slice, next = b.Slice(next)
   288  		if len(slice) == 0 {
   289  			continue
   290  		}
   291  		if err := f(slice); err != nil {
   292  			return err
   293  		}
   294  	}
   295  	return nil
   296  }
   297  
   298  const (
   299  	UseCalloc BufferType = iota
   300  	UseMmap
   301  	UseInvalid
   302  )
   303  
   304  type BufferType int
   305  
   306  func (t BufferType) String() string {
   307  	switch t {
   308  	case UseCalloc:
   309  		return "UseCalloc"
   310  	case UseMmap:
   311  		return "UseMmap"
   312  	default:
   313  		return "UseInvalid"
   314  	}
   315  }
   316  
   317  type (
   318  	LessFunc   func(a, b []byte) bool //nolint:revive //adopt fork, do not touch it
   319  	sortHelper struct {
   320  		offsets []int
   321  		b       *Buffer
   322  		tmp     *Buffer
   323  		less    LessFunc
   324  		small   []int
   325  	}
   326  )
   327  
   328  func (s *sortHelper) sortSmall(start, end int) {
   329  	s.tmp.Reset()
   330  	s.small = s.small[:0]
   331  	next := start
   332  	for next >= 0 && next < end {
   333  		s.small = append(s.small, next)
   334  		_, next = s.b.Slice(next)
   335  	}
   336  
   337  	// We are sorting the slices pointed to by s.small offsets, but only moving the offsets around.
   338  	sort.Slice(s.small, func(i, j int) bool {
   339  		left, _ := s.b.Slice(s.small[i])
   340  		right, _ := s.b.Slice(s.small[j])
   341  		return s.less(left, right)
   342  	})
   343  	// Now we iterate over the s.small offsets and copy over the slices. The result is now in order.
   344  	for _, off := range s.small {
   345  		s.tmp.Write(rawSlice(s.b.buf[off:]))
   346  	}
   347  	assert(end-start == copy(s.b.buf[start:end], s.tmp.Bytes()))
   348  }
   349  
   350  func assert(b bool) {
   351  	if !b {
   352  		log.Fatal("Assertion failure")
   353  	}
   354  }
   355  
   356  func check(err error) {
   357  	if err != nil {
   358  		log.Fatal(err.Error())
   359  	}
   360  }
   361  
   362  func check2(_ interface{}, err error) {
   363  	check(err)
   364  }
   365  
   366  func (s *sortHelper) merge(left, right []byte, start, end int) {
   367  	if len(left) == 0 || len(right) == 0 {
   368  		return
   369  	}
   370  	s.tmp.Reset()
   371  	check2(s.tmp.Write(left))
   372  	left = s.tmp.Bytes()
   373  
   374  	var ls, rs []byte
   375  
   376  	copyLeft := func() {
   377  		assert(len(ls) == copy(s.b.buf[start:], ls))
   378  		left = left[len(ls):]
   379  		start += len(ls)
   380  	}
   381  	copyRight := func() {
   382  		assert(len(rs) == copy(s.b.buf[start:], rs))
   383  		right = right[len(rs):]
   384  		start += len(rs)
   385  	}
   386  
   387  	for start < end {
   388  		if len(left) == 0 {
   389  			assert(len(right) == copy(s.b.buf[start:end], right))
   390  			return
   391  		}
   392  		if len(right) == 0 {
   393  			assert(len(left) == copy(s.b.buf[start:end], left))
   394  			return
   395  		}
   396  		ls = rawSlice(left)
   397  		rs = rawSlice(right)
   398  
   399  		// We skip the first 4 bytes in the rawSlice, because that stores the length.
   400  		if s.less(ls[4:], rs[4:]) {
   401  			copyLeft()
   402  		} else {
   403  			copyRight()
   404  		}
   405  	}
   406  }
   407  
   408  func (s *sortHelper) sort(lo, hi int) []byte {
   409  	assert(lo <= hi)
   410  
   411  	mid := lo + (hi-lo)/2
   412  	loff, hoff := s.offsets[lo], s.offsets[hi]
   413  	if lo == mid {
   414  		// No need to sort, just return the buffer.
   415  		return s.b.buf[loff:hoff]
   416  	}
   417  
   418  	// lo, mid would sort from [offset[lo], offset[mid]) .
   419  	left := s.sort(lo, mid)
   420  	// Typically we'd use mid+1, but here mid represents an offset in the buffer. Each offset
   421  	// contains a thousand entries. So, if we do mid+1, we'd skip over those entries.
   422  	right := s.sort(mid, hi)
   423  
   424  	s.merge(left, right, loff, hoff)
   425  	return s.b.buf[loff:hoff]
   426  }
   427  
   428  // SortSlice is like SortSliceBetween but sorting over the entire buffer.
   429  func (b *Buffer) SortSlice(less func(left, right []byte) bool) {
   430  	b.SortSliceBetween(b.StartOffset(), int(b.offset), less)
   431  }
   432  
   433  func (b *Buffer) SortSliceBetween(start, end int, less LessFunc) {
   434  	if start >= end {
   435  		return
   436  	}
   437  	if start == 0 {
   438  		panic("start can never be zero")
   439  	}
   440  
   441  	var offsets []int
   442  	next, count := start, 0
   443  	for next >= 0 && next < end {
   444  		if count%1024 == 0 {
   445  			offsets = append(offsets, next)
   446  		}
   447  		_, next = b.Slice(next)
   448  		count++
   449  	}
   450  	assert(len(offsets) > 0)
   451  	if offsets[len(offsets)-1] != end {
   452  		offsets = append(offsets, end)
   453  	}
   454  
   455  	szTmp := int(float64((end-start)/2) * 1.1)
   456  	s := &sortHelper{
   457  		offsets: offsets,
   458  		b:       b,
   459  		less:    less,
   460  		small:   make([]int, 0, 1024),
   461  		tmp:     NewBuffer(szTmp, b.tag),
   462  	}
   463  	defer s.tmp.Release()
   464  
   465  	left := offsets[0]
   466  	for _, off := range offsets[1:] {
   467  		s.sortSmall(left, off)
   468  		left = off
   469  	}
   470  	s.sort(0, len(offsets)-1)
   471  }
   472  
   473  func rawSlice(buf []byte) []byte {
   474  	sz := binary.BigEndian.Uint32(buf)
   475  	return buf[:4+int(sz)]
   476  }
   477  
   478  // Slice would return the slice written at offset.
   479  func (b *Buffer) Slice(offset int) ([]byte, int) {
   480  	if offset >= int(b.offset) {
   481  		return nil, -1
   482  	}
   483  
   484  	sz := binary.BigEndian.Uint32(b.buf[offset:])
   485  	start := offset + 4
   486  	next := start + int(sz)
   487  	res := b.buf[start:next]
   488  	if next >= int(b.offset) {
   489  		next = -1
   490  	}
   491  	return res, next
   492  }
   493  
   494  // SliceOffsets is an expensive function. Use sparingly.
   495  func (b *Buffer) SliceOffsets() []int {
   496  	next := b.StartOffset()
   497  	var offsets []int
   498  	for next >= 0 {
   499  		offsets = append(offsets, next)
   500  		_, next = b.Slice(next)
   501  	}
   502  	return offsets
   503  }
   504  
   505  func (b *Buffer) Data(offset int) []byte {
   506  	if offset > b.curSz {
   507  		panic("offset beyond current size")
   508  	}
   509  	return b.buf[offset:b.curSz]
   510  }
   511  
   512  // Write would write p bytes to the buffer.
   513  func (b *Buffer) Write(p []byte) (n int, err error) {
   514  	n = len(p)
   515  	b.Grow(n)
   516  	assert(n == copy(b.buf[b.offset:], p))
   517  	b.offset += uint64(n)
   518  	return n, nil
   519  }
   520  
   521  // Reset would reset the buffer to be reused.
   522  func (b *Buffer) Reset() {
   523  	b.offset = uint64(b.StartOffset())
   524  }
   525  
   526  // Release would free up the memory allocated by the buffer. Once the usage of buffer is done, it is
   527  // important to call Release, otherwise a memory leak can happen.
   528  func (b *Buffer) Release() error {
   529  	if b == nil {
   530  		return nil
   531  	}
   532  	switch b.bufType {
   533  	case UseCalloc:
   534  		Free(b.buf)
   535  	case UseMmap:
   536  		if b.mmapFile == nil {
   537  			return nil
   538  		}
   539  		path := b.mmapFile.Fd.Name()
   540  		if err := b.mmapFile.Close(-1); err != nil {
   541  			return fmt.Errorf("%w while closing file: %s", err, path)
   542  		}
   543  		if !b.persistent {
   544  			if err := os.Remove(path); err != nil {
   545  				return fmt.Errorf("%w while deleting file %s", err, path)
   546  			}
   547  		}
   548  	}
   549  	return nil
   550  }