github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/aq/db.go (about)

     1  package aq
     2  
     3  import (
     4  	"errors"
     5  	"os"
     6  	"sync"
     7  
     8  	"github.com/edsrzf/mmap-go"
     9  )
    10  
    11  type DB struct {
    12  	filename string
    13  	size     int32
    14  	file     *os.File
    15  
    16  	data mmap.MMap
    17  	mu   sync.Mutex
    18  	head int32
    19  }
    20  
    21  func New(filename string, size int) (*DB, error) {
    22  	db := &DB{filename: filename, size: int32(size)}
    23  	return db, db.open()
    24  }
    25  
    26  func (db *DB) open() error {
    27  	var err error
    28  	db.file, err = os.OpenFile(db.filename, os.O_RDWR|os.O_CREATE, 0777)
    29  	if err != nil {
    30  		return err
    31  	}
    32  	if err := db.file.Truncate(int64(db.size)); err != nil {
    33  		return err
    34  	}
    35  
    36  	db.data, err = mmap.Map(db.file, os.O_RDWR, 0)
    37  	return err
    38  }
    39  
    40  func (db *DB) Flush() error { return db.data.Flush() }
    41  
    42  func (db *DB) Close() error {
    43  	err := db.data.Unmap()
    44  	db.data = nil
    45  
    46  	err2 := db.file.Close()
    47  	db.file = nil
    48  
    49  	if err == nil {
    50  		return err2
    51  	}
    52  	return err
    53  }
    54  
    55  var (
    56  	InvalidHeader   = Header{Off: -1, Len: -1}
    57  	ErrFull         = errors.New("database is full")
    58  	ErrInvalidValue = errors.New("empty value is not allowed")
    59  )
    60  
    61  type Header struct {
    62  	Off int32
    63  	Len int32
    64  }
    65  
    66  func (db *DB) Write(data []byte) (Header, error) {
    67  	if len(data) == 0 {
    68  		return InvalidHeader, ErrInvalidValue
    69  	}
    70  	db.mu.Lock()
    71  	if db.head+4+int32(len(data)) > db.size {
    72  		db.mu.Unlock()
    73  		return InvalidHeader, ErrFull
    74  	}
    75  	hdr := Header{
    76  		Off: int32(db.head + 4),
    77  		Len: int32(len(data)),
    78  	}
    79  
    80  	db.data[db.head] = byte(hdr.Len >> 24)
    81  	db.data[db.head+1] = byte(hdr.Len >> 16)
    82  	db.data[db.head+2] = byte(hdr.Len >> 8)
    83  	db.data[db.head+3] = byte(hdr.Len >> 0)
    84  
    85  	db.head += 4 + hdr.Len
    86  	db.mu.Unlock()
    87  
    88  	copy(db.data[hdr.Off:hdr.Off+hdr.Len], data)
    89  	return hdr, nil
    90  }
    91  
    92  func (db *DB) Read(hdr Header) ([]byte, error) { return db.data[hdr.Off : hdr.Off+hdr.Len], nil }
    93  
    94  type Iterator struct {
    95  	head int32
    96  	clen int32
    97  	data []byte
    98  }
    99  
   100  func (db *DB) Iterate() Iterator { return Iterator{0, 0, db.data[:db.head]} }
   101  
   102  func (it *Iterator) Next() bool {
   103  	// jump to next header
   104  	it.head += it.clen
   105  	it.clen = 0
   106  
   107  	if it.head+4 > int32(len(it.data)) {
   108  		return false
   109  	}
   110  
   111  	it.clen = int32(it.data[it.head+0])<<24 |
   112  		int32(it.data[it.head+1])<<16 |
   113  		int32(it.data[it.head+2])<<8 |
   114  		int32(it.data[it.head+3])<<0
   115  	it.head += 4
   116  	return it.head+it.clen <= int32(len(it.data))
   117  }
   118  
   119  func (it *Iterator) Bytes() []byte { return it.data[it.head : it.head+it.clen] }