Implementing a key-value database in Golang, Part 4 - Merging and Compaction

type index struct {
	table  map[string]int64
	cursor int64
}

func (i *index) get(key string) (int64, bool) {
	val, ok := i.table[key]
	return val, ok
}

func (i *index) put(key string, written int64) {
	i.table[key] = i.cursor
	i.cursor += written
}

type segment struct {
	storagePath   string
	maxRecordSize int
	logger        *log.Logger
	async         bool
	suffix        string

	index *index
}

func newSegment(baseDir string, maxRecordSize int, async bool, logger *log.Logger) *segment {
	filename := genFileName(openSegmentSuffix)
	filePath := path.Join(baseDir, filename)

	return &segment{
		storagePath: filePath,
		logger:      logger,
		async:       async,
		index: &index{
			cursor: 0,
			table:  map[string]int64{},
		},
	}
}

func fromFile(filePath string, maxRecordSize int, async bool, logger *log.Logger) (*segment, error) {
	idx := index{
		cursor: 0,
		table:  map[string]int64{},
	}

	f, err := os.OpenFile(filePath, os.O_RDONLY|os.O_CREATE, 0600)
	defer f.Close()
	if err != nil {
		return nil, err
	}

	scanner, err := record.NewScanner(f, maxRecordSize)
	if err != nil {
		return nil, err
	}

	for scanner.Scan() {
		record := scanner.Record()
		idx.put(record.Key(), int64(record.Size()))
	}

	if scanner.Err() != nil {
		return nil, fmt.Errorf("could not scan entry, %w", err)
	}

	return &segment{
		storagePath: filePath,
		index:       &idx,
	}, nil
}

func (s *segment) get(key string) (*record.Record, error) {
	offset, ok := s.index.get(key)
	if !ok {
		return nil, kvdb.NewNotFoundError(key)
	}

	f, err := s.getFile(os.O_RDONLY)
	defer f.Close()
	if err != nil {
		return nil, err
	}

	_, err = f.Seek(offset, io.SeekStart)
	if err != nil {
		return nil, err
	}

	scanner, err := record.NewScanner(f, s.maxRecordSize)
	if err != nil {
		return nil, err
	}

	if scanner.Scan() {
		return scanner.Record(), nil
	}

	return nil, kvdb.NewNotFoundError(key)
}

func (s *segment) append(record *record.Record) error {
	file, err := s.getFile(os.O_CREATE | os.O_WRONLY | os.O_APPEND)
	defer file.Close()
	if err != nil {
		return fmt.Errorf("could not open file: %s for write, %w", s.storagePath, err)
	}

	n, err := record.Write(file)
	if err != nil {
		return fmt.Errorf("could not write record to file: %s, %w", s.storagePath, err)
	}

	if !s.async {
		if err := file.Sync(); err != nil {
			return err
		}
	}

	if err := file.Close(); err != nil {
		return err
	}

	s.index.put(record.Key(), int64(n))
	return nil
}

func (s *segment) getFile(mode int) (*os.File, error) {
	return os.OpenFile(s.storagePath, mode, 0600)
}

func (s *segment) changeSuffix(oldSuffix, newSuffix string) error {
	newFilePath := strings.Replace(s.storagePath, oldSuffix, newSuffix, 1)

	if err := os.Rename(s.storagePath, newFilePath); err != nil {
		return err
	}

	s.storagePath = newFilePath
	return nil
}

func (s *segment) size() int64 {
	return s.index.cursor
}

func (s *segment) clearFile() error {
	s.index.table = map[string]int64{}
	s.index.cursor = 0
	return os.Remove(s.storagePath)
}

type segmentStack struct {
	segments []*segment
}

func newSegmentStack() *segmentStack {
	segments := make([]*segment, 0)
	return &segmentStack{
		segments: segments,
	}
}

func (s *segmentStack) iter() *segmentStackIter {
	return &segmentStackIter{
		segments: s.segments,
		pos:      -1,
	}
}

func (s *segmentStack) push(seg *segment) {
	s.segments = append([]*segment{seg}, s.segments...)
}

type segmentStackIter struct {
	segments []*segment
	pos      int
}

func (i *segmentStackIter) hasNext() bool {
	return i.pos < len(i.segments)-1
}

func (i *segmentStackIter) next() *segment {
	i.pos = i.pos + 1
	return i.segments[i.pos]
}

const (
	defaultAsync          = false
	defaultMaxRecordSize  = 1024 * 1024 //1Mb
	defaultSegmentMaxSize = 4096 * 1024 //4Mb
	closedSegmentSuffix   = ".cseg"
	openSegmentSuffix     = ".oseg"
)

var voidLogger = log.New(ioutil.Discard, "", log.LstdFlags)

// Store implements the kvdb store interface providing a simple key-value
// database engine based on an append-log.
type Store struct {
	storagePath    string
	maxRecordSize  int
	maxSegmentSize int64
	logger         *log.Logger
	async          bool

	openSegment *segment

	// closedSegmentStack contains all immutable segments. The newest segment
	// has position 0 and the oldest at the last position
	closedSegments *segmentStack
}

// Config contains the configuration properties for the compacted aol store
type Config struct {
	// Storage path
	BasePath string
	// Sets a limit on the size of the records
	MaxRecordSize *int
	// If true, fsync will be called on every write
	Async *bool
	// MaxSegmentSize defines the maximum size of the segments, must be >=
	// MaxRecordSize
	MaxSegmentSize *int

	Logger *log.Logger
}

// NewStore returns a new SimpleLogStore
func NewStore(config Config) (*Store, error) {
	var (
		maxRecordSize  = defaultMaxRecordSize
		maxSegmentSize = defaultSegmentMaxSize
		storagePath    = config.BasePath
		async          = defaultAsync
		logger         = voidLogger
	)

	if _, err := os.OpenFile(storagePath, os.O_CREATE, 0600); err != nil {
		return nil, err
	}

	if config.MaxRecordSize != nil {
		maxRecordSize = *config.MaxRecordSize
	}

	if config.MaxSegmentSize != nil {
		maxSegmentSize = *config.MaxSegmentSize
	}

	if config.Async != nil {
		async = *config.Async
	}

	if config.Logger != nil {
		logger = config.Logger
	}

	openSegment, err := loadOpenSegment(storagePath, maxRecordSize, async, logger)
	if err != nil {
		return nil, err
	}

	closedSegments, err := loadClosedSegments(storagePath, maxRecordSize, async, logger)
	if err != nil {
		return nil, err
	}

	store := &Store{
		storagePath:    storagePath,
		maxRecordSize:  maxRecordSize,
		maxSegmentSize: int64(maxSegmentSize),
		async:          async,
		logger:         logger,
		openSegment:    openSegment,
		closedSegments: closedSegments,
	}

	return store, nil
}

func loadOpenSegment(storagePath string, maxRecordSize int, async bool, log *log.Logger) (*segment, error) {
	fi, err := listFilesWithSuffix(storagePath, openSegmentSuffix, true)
	if err != nil {
		return nil, fmt.Errorf("could not load open segment: %w", err)
	}

	switch len(fi) {
	case 0:
		return newSegment(storagePath, maxRecordSize, async, log), nil
	case 1:
		return fromFile(fi[0].filepath, maxRecordSize, async, log)
	default:
		return nil, fmt.Errorf("more than one open segment found")
	}
}

func loadClosedSegments(storagePath string, maxRecordSize int, async bool, log *log.Logger) (*segmentStack, error) {
	fis, err := listFilesWithSuffix(storagePath, closedSegmentSuffix, true)
	if err != nil {
		return nil, fmt.Errorf("could not list closed segment files: %w", err)
	}

	closedSegments := newSegmentStack()
	for _, fi := range fis {
		s, err := fromFile(fi.filepath, maxRecordSize, async, log)
		if err != nil {
			return nil, fmt.Errorf("could not load closed segment: %s, %w", fi.filepath, err)
		}
		closedSegments.push(s)
	}

	return closedSegments, nil
}

// Close closes the store
func (s *Store) Close() error {
	s.logger.Println("Closing database")
	return nil
}

// IsNotFoundError returns true if the error signales that a non-existing key
// was requested
func (s *Store) IsNotFoundError(err error) bool {
	return kvdb.IsNotFoundError(err)
}

// IsBadRequestError returns true if the error, or any of the wrapped errors
// is of type BadRequestError
func (s *Store) IsBadRequestError(err error) bool {
	return kvdb.IsBadRequestError(err)
}

// Get returns the value associated with the key or a kvdb.NotFoundError if the
// key was not found, or any other error encountered
func (s *Store) Get(key string) ([]byte, error) {
	record, err := s.openSegment.get(key)

	if err == nil {
		return resolveRecord(record)
	} else if !kvdb.IsNotFoundError(err) {
		return nil, err
	}

	iter := s.closedSegments.iter()

	for iter.hasNext() {
		record, err := iter.next().get(key)
		if err == nil {
			return resolveRecord(record)
		} else if !kvdb.IsNotFoundError(err) {
			return nil, err
		}
	}

	return nil, kvdb.NewNotFoundError(key)
}

func resolveRecord(record *record.Record) ([]byte, error) {
	if record.IsTombstone() {
		return nil, kvdb.NewNotFoundError(record.Key())
	}

	return record.Value(), nil
}

// Put saves the value to the database and returns any error encountered
func (s *Store) Put(key string, value []byte) error {
	record := record.NewValue(key, value)

	if record.Size() > s.maxRecordSize {
		msg := fmt.Sprintf("key-value too big, max size: %d", s.maxRecordSize)
		return kvdb.NewBadRequestError(msg)
	}

	return s.append(record)
}

// Delete removes the value from the store and returns any error encountered
func (s *Store) Delete(key string) error {
	record := record.NewTombstone(key)
	return s.append(record)
}

func (s *Store) append(record *record.Record) error {
	if s.openSegment.size()+int64(record.Size()) > s.maxSegmentSize {
		if err := s.rotateOpenSegment(); err != nil {
			return err
		}
	}

	return s.openSegment.append(record)
}

func (s *Store) rotateOpenSegment() error {
	newOpenSegment := newSegment(s.storagePath, s.maxRecordSize, s.async, s.logger)
	if err := s.openSegment.changeSuffix(openSegmentSuffix, closedSegmentSuffix); err != nil {
		return err
	}

	s.closedSegments.push(s.openSegment)
	s.openSegment = newOpenSegment
}

const (
	readers     = 20
	writers     = 40
	removers    = 2
	nWrites     = 1000
	nReads      = 1000
	nRemoves    = 1000
	cardinality = 1000
)

var data string = `{
  "extends": [
    "tslint:latest",
    "tslint-config-standard",
    "tslint-react"
  ],
  "rules": {
    "semicolon": [true, "never"],
    "object-literal-sort-keys": false,
    "trailing-comma": [true, {"multiline": "always", "singleline": "never"}],
    "jsx-no-lambda": false,
    "jsx-no-multiline-js": false,
    "quotemark": [true, "single", "jsx-double"],
    "no-implicit-dependencies": [true, "dev"],
    "no-console": [false],
    "max-line-length": [true, 200]
  }
}`

func TestLoad(t *testing.T) {
	dbPath, err := ioutil.TempDir("./", dbPathPattern)
	if err != nil {
		t.Fatal(err)
	}

	defer os.RemoveAll(dbPath)

	logger := log.New(os.Stdout, "", log.LstdFlags)
	maxRecordSize := 5 * 1024
	maxSegmentSize := 20 * 1024
	async := true

	store, err := compactedaol.NewStore(compactedaol.Config{
		Async:          &async,
		Logger:         logger,
		BasePath:       dbPath,
		MaxRecordSize:  &maxRecordSize,
		MaxSegmentSize: &maxSegmentSize,
	})

	if err != nil {
		t.Fatal(err)
	}

	run(t, store)
}

func run(t *testing.T, store kvdb.Store) {
	wg := sync.WaitGroup{}
	for i := 0; i < writers; i++ {
		wg.Add(1)
		go func(wi int, gr *sync.WaitGroup) {
			writer(t, store)
			gr.Done()
		}(i, &wg)
	}

	for i := 0; i < readers; i++ {
		wg.Add(1)
		go func(ri int, gr *sync.WaitGroup) {
			reader(t, store)
			gr.Done()
		}(i, &wg)
	}

	for i := 0; i < removers; i++ {
		wg.Add(1)
		go func(ri int, gr *sync.WaitGroup) {
			remover(t, store)
			gr.Done()
		}(i, &wg)
	}

	wg.Wait()
	store.Close()
}

func reader(t *testing.T, store kvdb.Store) {
	for i := 0; i < nReads; i++ {
		rnd := rand.Intn(cardinality)
		_, err := store.Get(fmt.Sprintf("key-%d", rnd))
		if err != nil && !kvdb.IsNotFoundError(err) {
			t.Fatal(err)
		}
	}
}

func writer(t *testing.T, store kvdb.Store) {
	for i := 0; i < nWrites; i++ {
		rnd := rand.Intn(cardinality)
		err := store.Put(fmt.Sprintf("key-%d", rnd), []byte(data))
		if err != nil {
			t.Fatal(err)
		}
	}
}

func remover(t *testing.T, store kvdb.Store) {
	for i := 0; i < nRemoves; i++ {
		rnd := rand.Intn(cardinality)
		err := store.Delete(fmt.Sprintf("key-%d", rnd))
		if err != nil {
			t.Fatal(err)
		}
	}
}

type index struct {
	table  map[string]int64
	mutex  sync.RWMutex
	cursor int64
}

func (i *index) get(key string) (int64, bool) {
	i.mutex.RLock()
	defer i.mutex.RUnlock()
	val, ok := i.table[key]
	return val, ok
}

func (i *index) put(key string, written int64) {
	i.mutex.Lock()
	defer i.mutex.Unlock()
	i.table[key] = i.cursor
	i.cursor += written
}

type onCompactionDoneFunc func(targetFile string, compactedFiles []string) error

type compacter struct {
	onCompactionDone    onCompactionDoneFunc
	compactionThreshold int
	storagePath         string
	maxRecordSize       int
	semaphore           *semaphore.Weighted
	logger              *log.Logger
}

func newCompacter(
	onCompactionDone onCompactionDoneFunc,
	compactionThreshold int,
	storagePath string,
	maxRecordSize int,
	logger *log.Logger) *compacter {

	return &compacter{
		onCompactionDone:    onCompactionDone,
		compactionThreshold: compactionThreshold,
		storagePath:         storagePath,
		maxRecordSize:       maxRecordSize,
		semaphore:           semaphore.NewWeighted(1),
		logger:              logger,
	}
}

func (c *compacter) run(interval time.Duration) {
	go (func() {
		for range time.Tick(interval) {
			go func() {
				if c.semaphore.TryAcquire(1) {
					defer c.semaphore.Release(1)
					c.logger.Println("compaction: running compaction")

					filesToCompact := c.selectFilesForCompaction()
					if filesToCompact == nil || len(filesToCompact) == 0 {
						c.logger.Println("compaction: no files to compact")
						return
					}

					targetFile := strings.ReplaceAll(filesToCompact[0], closedSegmentSuffix, compactionSuffix)
					if err := c.doCompaction(targetFile, filesToCompact); err != nil {
						c.logger.Printf("compaction: could not run compaction: %s", err)
						return
					}

					if err := c.onCompactionDone(targetFile, filesToCompact); err != nil {
						c.logger.Printf("compaction: could not run onCompactionDone: %s", err)
					}
				}
			}()
		}
	})()
}

func (c *compacter) stop() {
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
	defer cancel()

	c.semaphore.Acquire(ctx, 1)
	c.ticker.Stop()
	c.semaphore.Release(1)
}

// selectFilesForCompaction tries to select as many consecutive files as
// possible without exceeding the CompactionThreshold.
func (c *compacter) selectFilesForCompaction() []string {
	files, _ := listFilesWithSuffix(c.storagePath, closedSegmentSuffix, false)

	fileGroup := []string{}
	var groupSize int64 = 0
	for i := range files {
		fs := files[i].info
		if groupSize+fs.Size() < int64(c.compactionThreshold) {
			groupSize += fs.Size()
			fileGroup = append(fileGroup, files[i].filepath)
		} else {
			if len(fileGroup) > 1 {
				break
			} else {
				fileGroup = []string{files[i].filepath}
				groupSize = fs.Size()
			}
		}
	}

	if len(fileGroup) >= 2 {
		sort.Sort(sort.Reverse(sort.StringSlice(fileGroup)))
		return fileGroup
	}

	return nil
}

func (c *compacter) doCompaction(target string, sources []string) error {
	var sourceScanners = []*record.Scanner{}

	for _, source := range sources {
		file, err := os.Open(source)
		if err != nil {
			return fmt.Errorf("could not open file for compaction: %w", err)
		}
		defer file.Close()

		scanner, err := record.NewScanner(file, c.maxRecordSize)
		if err != nil {
			return err
		}

		sourceScanners = append([]*record.Scanner{scanner}, sourceScanners...)
	}

	data := map[string]*record.Record{}
	for _, source := range sourceScanners {
		for source.Scan() {
			if source.Err() != nil {
				c.logger.Fatalf("compacter: skipping corrupted record for key: %s\n", source.Record().Key())
			} else {
				record := source.Record()
				data[record.Key()] = record
			}
		}
	}

	c.logger.Printf("compacter: creating new segment: %s\n", target)
	targetFile, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
	if err != nil {
		return fmt.Errorf("could not create compaction target: %w", err)
	}
	defer targetFile.Close()

	for _, record := range data {
		_, err := record.Write(targetFile)
		if err != nil {
			return fmt.Errorf("compacter: could not write to target: %w", err)
		}
	}

	return nil
}

func (s *Store) onCompactionDone(targetFile string, compactedFiles []string) error {
	newSegment, err := fromFile(targetFile, s.maxRecordSize, s.async, s.logger)
	if err != nil {
		return fmt.Errorf("could not create segment of compaction target: %w", err)
	}

	err = s.closedSegments.replace(func(segment *segment) bool {
		return segment.storagePath == compactedFiles[0]
	}, newSegment)

	if err != nil {
		return fmt.Errorf("could not replace with compacted segment: %w", err)
	}

	filesToRemove := compactedFiles[1:]
	for i := range filesToRemove {
		err = s.closedSegments.remove(func(segment *segment) bool {
			return filesToRemove[i] == segment.storagePath
		})

		if err != nil {
			return fmt.Errorf("could not remove compacted segment: %w", err)
		}
	}

	if err = newSegment.changeSuffix(compactionSuffix, closedSegmentSuffix); err != nil {
		return fmt.Errorf("could not rename compacted segment: %w", err)
	}

	return nil
}

func (s *segmentStack) remove(predicate func(segment *segment) bool) error {
	rem := []*segment{}

	for i := range s.segments {
		if predicate(s.segments[i]) {
			if err := s.segments[i].clearFile(); err != nil {
				return fmt.Errorf("could not remove segment file: %s, due to: %w",
					s.segments[i].storagePath, err)
			}
		} else {
			rem = append(rem, s.segments[i])
		}
	}
	s.segments = rem

	return nil
}

func (s *segmentStack) replace(predicate func(segment *segment) bool, seg *segment) error {
	rem := []*segment{}

	for i := range s.segments {
		if predicate(s.segments[i]) {
			if err := s.segments[i].clearFile(); err != nil {
				return fmt.Errorf("could not remove segment file: %s, due to: %w",
					s.segments[i].storagePath, err)
			}
			rem = append(rem, seg)
		} else {
			rem = append(rem, s.segments[i])
		}
	}
	s.segments = rem

	return nil
}