github.com/blong14/gache@v0.0.0-20240124023949-89416fd8bbfa/internal/db/sstable/arena/allocator.go (about)

     1  package arena
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"errors"
     7  	gfile "github.com/blong14/gache/internal/io/file"
     8  	"os"
     9  	"runtime"
    10  	"sync"
    11  	"time"
    12  )
    13  
    14  // Arena https://gist.github.com/quillaja/222c9af7ade058b60ed08e13bf0b6387
    15  type Arena interface {
    16  	Free()
    17  	Bytes() []byte
    18  	Write([]byte) (int, error)
    19  	Append([]byte, []byte) (int, int, error)
    20  	Read([]byte) (int, error)
    21  	ReadAt([]byte, int64, int64) (int, error)
    22  	Cap() uint64
    23  	Len() uint64
    24  }
    25  
    26  type goarena struct {
    27  	data    gfile.Map
    28  	buffer  chan KeyValue
    29  	end     uint64
    30  	invalid bool
    31  }
    32  
    33  func flusher(a *goarena) {
    34  	ticker := time.NewTicker(100 * time.Millisecond)
    35  	buf := make([]byte, 0)
    36  	buffer := bytes.NewBuffer(buf)
    37  	for {
    38  		select {
    39  		case dat := <-a.buffer:
    40  			buffer.Write(dat.Key)
    41  			buffer.Write([]byte("::"))
    42  			buffer.Write(dat.Value)
    43  			buffer.Write([]byte(";"))
    44  		case <-ticker.C:
    45  			if buffer.Len() > 0 {
    46  				_, _ = a.Write(buffer.Bytes())
    47  				buffer.Reset()
    48  			}
    49  		}
    50  	}
    51  }
    52  
    53  func NewGoArena(f *os.File, size uint64) Arena {
    54  	m, err := gfile.NewMap(f, gfile.Prot(gfile.Read), gfile.Prot(gfile.Write), gfile.Flag(gfile.Shared))
    55  	if err != nil {
    56  		panic(err)
    57  	}
    58  	a := &goarena{
    59  		data:    m,
    60  		buffer:  make(chan KeyValue, 4096),
    61  		end:     0,
    62  		invalid: false,
    63  	}
    64  	go flusher(a)
    65  
    66  	return a
    67  }
    68  
    69  func (a *goarena) Free() {
    70  	err := a.data.Close()
    71  	if err != nil {
    72  		panic(err)
    73  	}
    74  	a.buffer = nil
    75  	a.invalid = true
    76  	a.end = 0
    77  	runtime.GC()
    78  }
    79  
    80  func (a *goarena) Write(p []byte) (int, error) {
    81  	return a.data.Write(p)
    82  }
    83  
    84  type KeyValue struct {
    85  	Key   json.RawMessage `json:"key"`
    86  	Value json.RawMessage `json:"value"`
    87  }
    88  
    89  var pool = sync.Pool{New: func() interface{} {
    90  	return KeyValue{}
    91  }}
    92  
    93  func (a *goarena) Append(k, v []byte) (int, int, error) {
    94  	var out []byte
    95  	buf := bytes.NewBuffer(out)
    96  	buf.Write(k)
    97  	buf.Write([]byte("::"))
    98  	buf.Write(v)
    99  	buf.Write([]byte(";"))
   100  	return a.data.Append(buf.Bytes())
   101  }
   102  
   103  func (a *goarena) Read(p []byte) (int, error) {
   104  	return a.data.Read(p)
   105  }
   106  
   107  func (a *goarena) ReadAt(p []byte, start, len_ int64) (int, error) {
   108  	kv := make([]byte, len_)
   109  	_, err := a.data.Peek(kv, start, len_)
   110  	if err != nil {
   111  		return -1, err
   112  	}
   113  	values := bytes.Split(kv, []byte("::"))
   114  	if len(values) != 2 {
   115  		return -1, errors.New("malformatted kv pair")
   116  	}
   117  	v := bytes.TrimSuffix(values[1], []byte(";"))
   118  	n := copy(p, v)
   119  	return n, nil
   120  }
   121  
   122  func (a *goarena) Bytes() []byte {
   123  	return a.data.Bytes()
   124  }
   125  
   126  func (a *goarena) Cap() uint64 { return uint64(a.data.Len()) }
   127  func (a *goarena) Len() uint64 { return uint64(a.data.Pos()) }