Olivine: Building a Redis-Compatible Service in Go. · 章節目錄

第 13 章

[Day 13] Key Expiration

We stored data into map[string]string now.

type mapStorage struct {
	storage map[string]string
	mu      sync.RWMutex
}

Obviously, it cannot handle data expiration feature by this simple data structure directly.

How to implement?

Redis keys are expired in two ways: a passive way and active way.

A key is passively expired when a client tries to access it and the key is timed out.

However, this is not enough as there are expired keys that will never be accessed again. These keys should be expired anyway, so periodically, Redis tests a few keys at random amongst the set of keys with an expiration. All the keys that are already expired are deleted from the keyspace.

Ref

Therefore, we need to store not only string but also expiration for values.

The Object{}

“Object” is used to be a collection of data, we defined the Object{} type is a value for storage to replace string.

package object

import "time"

type Object interface {
	Key() string
	ExpiresAt() *time.Time
}

type base struct {
	key       string
	expiresAt *time.Time
}

func (obj base) Key() string {
	return obj.key
}

func (obj base) ExpiresAt() *time.Time {
	return obj.expiresAt
}

“String” is an Object{} type which can be stored.

package object

import "time"

type String struct {
	base

	val string
}

func NewString(k, v string, expiresAt *time.Time) *String {
	return &String{
		val:  v,
		base: base{key: k, expiresAt: expiresAt},
	}
}

func (str *String) String() string {
	return str.val
}

The mapStorage{} should be updated:

// file: internal/repo/storage.go
func NewStorage() Storage {
	s := mapStorage{}
-	s.storage = make(map[string]string)
+	s.storage = make(map[string]object.Object)

	return &s
}

type mapStorage struct {
-	storage map[string]string
+	storage map[string]object.Object
	mu      sync.Mutex
}
func (s *mapStorage) Set(_ context.Context, obj object.Object) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.storage[obj.Key()] = obj
	return nil
}

func (s *mapStorage) Get(_ context.Context, k string) (v object.Object, err error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var ok bool
	if v, ok = s.storage[k]; !ok {
		return nil, fmt.Errorf("%w: miss", ErrNotFound)
	}
	if v.ExpiresAt() != nil && time.Now().After(*v.ExpiresAt()) {
		delete(s.storage, k) // remove key from storage when expired
		return nil, fmt.Errorf("%w: expired", ErrNotFound)
	}

	return
}

It removes expired key from map when getting, it is passive deletion.

Make SET supports expiration

The SET command can set a key with expiration by:

  • EX {sec}: data expires after {sec} seconds.
  • PX {msec}: data expires after {msec} milliseconds.
  • EXAT {usec}: data expires at {usec} unix epoch in seconds.
  • PXAT {umsec}: data expires at {umsec} unix epoch in milliseconds.

The full arguments format of SET command:

SET key value [NX | XX | IFEQ ifeq-value | IFNE ifne-value |
  IFDEQ ifdeq-digest | IFDNE ifdne-digest] [GET] [EX seconds |
  PX milliseconds | EXAT unix-time-seconds |
  PXAT unix-time-milliseconds | KEEPTTL]

Make a parse() for Set{} with returning (parsedSet, error):

func (c *Set) parse(cmd *resp.Command) (p parsedSet, err error) {
	const (
		awaitingKey = iota
		awaitingValue
		awaitingOption
		awaitingConditionValue
		awaitingPostConditionOption
		awaitingPostGetOption
		awaitingExpirationValue
		done
		invalid
	)

	state := awaitingKey
	var conditionType int
	var expirationOption string
	var expirationOptionIndex int

	for i, arg := range cmd.Args() {
		switch state {
		case awaitingKey:
			p.K = arg.String()
			state = awaitingValue
		case awaitingValue:
			p.V = arg.String()
			state = awaitingOption
		case awaitingOption:
			switch strings.ToUpper(arg.String()) {
			case "NX":
				p.Cond.Typ = nx
				state = awaitingPostConditionOption
			case "XX":
				p.Cond.Typ = xx
				state = awaitingPostConditionOption
			case "IFEQ":
				conditionType = ifeq
				state = awaitingConditionValue
			case "IFNE":
				conditionType = ifne
				state = awaitingConditionValue
			case "IFDEQ":
				conditionType = ifdeq
				state = awaitingConditionValue
			case "IFDNE":
				conditionType = ifdne
				state = awaitingConditionValue
			case "GET":
				p.Get = true
				state = awaitingPostGetOption
			case "EX", "PX", "EXAT", "PXAT":
				expirationOption = strings.ToUpper(arg.String())
				expirationOptionIndex = i
				state = awaitingExpirationValue
			case "KEEPTTL":
				p.Exp = new(time.Time)
				state = done
			default:
				state = invalid
			}
		case awaitingConditionValue:
			p.Cond = parseSetCond{Typ: conditionType, Val: arg.String()}
			state = awaitingPostConditionOption
		case awaitingPostConditionOption:
			switch strings.ToUpper(arg.String()) {
			case "GET":
				p.Get = true
				state = awaitingPostGetOption
			case "EX", "PX", "EXAT", "PXAT":
				expirationOption = strings.ToUpper(arg.String())
				expirationOptionIndex = i
				state = awaitingExpirationValue
			case "KEEPTTL":
				p.Exp = new(time.Time)
				state = done
			default:
				state = invalid
			}
		case awaitingPostGetOption:
			switch strings.ToUpper(arg.String()) {
			case "EX", "PX", "EXAT", "PXAT":
				expirationOption = strings.ToUpper(arg.String())
				expirationOptionIndex = i
				state = awaitingExpirationValue
			case "KEEPTTL":
				p.Exp = new(time.Time)
				state = done
			default:
				state = invalid
			}
		case awaitingExpirationValue:
			p.Exp, err = parseExpiration(expirationOption, arg.String())
			if err != nil {
				return parsedSet{}, fmt.Errorf("%w: %w", ErrSyntax, err)
			}
			state = done
		case done:
			state = invalid
		}
	}

	if state != awaitingOption && state != awaitingPostConditionOption && state != awaitingPostGetOption && state != done {
		return parsedSet{}, fmt.Errorf("%w: invalid SET syntax", ErrSyntax)
	}

	if p.Exp != nil && !p.Exp.IsZero() {
		cmd.UpdateAOF(expirationOptionIndex+1, resp.NewBulkString("PXAT"))
		cmd.UpdateAOF(expirationOptionIndex+2, resp.NewBulkString(strconv.FormatInt(p.Exp.UnixMilli(), 10)))
	}

	return p, nil
}

func parseExpiration(option, value string) (*time.Time, error) {
	switch option {
	case "EX", "PX":
		unit := "s"
		if option == "PX" {
			unit = "ms"
		}

		duration, err := parseDuration(value, unit)
		if err != nil {
			return nil, err
		}

		return new(time.Now().Add(duration)), nil
	case "EXAT", "PXAT":
		timestamp, err := strconv.Atoi(value)
		if err != nil {
			return nil, err
		}
		if timestamp <= 0 {
			return nil, errors.New("invalid expire time")
		}
		if option == "EXAT" {
			return new(time.Unix(int64(timestamp), 0)), nil
		}

		return new(time.UnixMilli(int64(timestamp))), nil
	}

	return nil, fmt.Errorf("unknown expiration option %q", option)
}

func parseDuration(s, unit string) (time.Duration, error) {
	if _, err := strconv.Atoi(s); err != nil {
		return 0, fmt.Errorf("invalid expire time: %w", err)
	}

	var sb strings.Builder
	sb.WriteString(s)
	sb.WriteString(unit)

	if d, err := time.ParseDuration(sb.String()); err != nil {
		return d, err
	} else if d <= 0 {
		return d, errors.New("invalid expire time")
	} else {
		return d, nil
	}
}

const (
	nx    = iota + 1 // set value only not exists
	xx               // set value only exists
	ifeq             // set only value == cond.Val
	ifne             // set only value != cond.Val
	ifdeq            // set only XXH3(value) == cond.Val
	ifdne            // set only XXH3(value) != cond.Val
)

type parseSetCond struct {
	Typ int
	Val string
}

type parsedSet struct {
	K    string
	V    string
	Cond parseSetCond
	Get  bool
	Exp  *time.Time
}

Note: I built a parser with a finiate state machine, it makes code very long but clear for state transfering.

SetParam{} interface

Because we parse() in Set{} command, and returns parsedSet{} as result instead of put object.Object as parameter for repo.Storage.Set().

It’s a good design for recieving repo.SetParam{} interface, and handle parsedSet{} when calls repo.Storage.Set()

type SetParam interface {
	Obj() object.Object
}
type parsedSet struct {
	K    string
	V    string
	Cond parseSetCond
	Get  bool
	Exp  *time.Time
}

var _ repo.SetParam = parsedSet{}

func (parsed parsedSet) Obj() object.Object {
	return object.NewString(parsed.K, parsed.V, parsed.Exp)
}

Use var _ repo.SetParam = parsedSet{} to ensure parsedSet{} is a repo.SetParam{} while compiling.

It’s time to update repo.Storage.Set() function:

func (s *mapStorage) Set(_ context.Context, param SetParam) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.storage[param.Obj().Key()] = param.Obj()
	return nil
}

The reason of using SetParam{} to handle parsedSet{} can make repo.Storage.Set() handle not only object.String{} but object.Hash{} in future for HSET Redis command.

Active way expiration

We’ve implemented passive way for removing key from storage if expired when accessing.

How’s the active way?

Redis spawn a thread, it cyclically get some keys from storage and remove it if expired every 100 millisecons.

Refactor with Worker{}

We’ve implemented a thread like feature: AOF syncing every second (if configured), although we implemented it by goroutine.

I decided to create a service.Worker{} to manage these gorutines: more clear flow in Application.Run():

package service

import (
	"context"
	"errors"
	"time"

	"golang.org/x/sync/errgroup"

	"olivine/internal/data"
	"olivine/internal/repo"
)

type Worker interface {
	Start(context.Context) error
}

func NewWorker(cfg *data.Config, aof AOF, storage repo.Storage) Worker {
	return &worker{
		cfg: cfg,

		aof:     aof,
		storage: storage,
	}
}

type worker struct {
	cfg *data.Config

	storage repo.Storage
	aof     AOF
}

func (w *worker) Start(ctx context.Context) error {
	g, ctx := errgroup.WithContext(ctx)

	if w.cfg.AOFEnabled && w.cfg.AOFFsync == data.AOFFsyncEverySec {
		g.Go(w.aofSyncer(ctx))
	}

	g.Go(w.storagePruner(ctx))

	return g.Wait()
}

func (w *worker) aofSyncer(ctx context.Context) func() error {
	return func() error {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return nil
			case <-ticker.C:
				if err := w.aof.Sync(); err != nil {
					return err
				}
			}
		}
	}
}

And call Worker.Start() in Application.Run() errgroup goroutine:

package service

import (
	"context"
	"time"

	"golang.org/x/sync/errgroup"

	"olivine/internal/data"
)

type Worker interface {
	Start(context.Context) error
}

func NewWorker(cfg *data.Config, aof AOF) Worker {
	return &worker{
		cfg: cfg,
		aof: aof,
	}
}

type worker struct {
	cfg *data.Config

	aof AOF
}

func (w *worker) Start(ctx context.Context) error {
	g, ctx := errgroup.WithContext(ctx)

	if w.cfg.AOFEnabled && w.cfg.AOFFsync == data.AOFFsyncAlways {
		g.Go(w.aofSyncer(ctx))
	}

	return g.Wait()
}

func (w *worker) aofSyncer(ctx context.Context) func() error {
	return func() error {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return nil
			case <-ticker.C:
				if err := w.aof.Sync(); err != nil {
					return err
				}
			}
		}
	}
}

Storage.Prune() and register in Worker{}

func (s *mapStorage) Prune(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			if stop := s.tryPrune(); stop {
				return nil
			}
		}
	}
}

func (s *mapStorage) tryPrune() bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	if len(s.storage) == 0 {
		return true
	}

	const sampleSize = 10

	sampled := 0
	expired := 0
	now := time.Now()

	for k, v := range s.storage {
		if sampled == sampleSize {
			break
		}

		sampled++
		if v.ExpiresAt() != nil && now.After(*v.ExpiresAt()) {
			delete(s.storage, k)
			expired++
		}
	}

	return expired*4 < sampled
}
// file: internal/service/worker.go
func (w *worker) Start(ctx context.Context) error {
	g, ctx := errgroup.WithContext(ctx)

	if w.cfg.AOFEnabled && w.cfg.AOFFsync == data.AOFFsyncEverySec {
		g.Go(w.aofSyncer(ctx))
	}

+	g.Go(w.storagePruner(ctx))

	return g.Wait()
}

+func (w *worker) storagePruner(ctx context.Context) func() error {
+	return func() error {
+		ticker := time.NewTicker(time.Millisecond * 100)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return nil
+			case <-ticker.C:
+				if err := w.storage.Prune(ctx); err != nil {
+					if ctx.Err() != nil && errors.Is(err, ctx.Err()) {
+						return nil
+					}
+					return err
+				}
+			}
+		}
+	}
+}

AOF for expiration keys

We stores keys command directly when appending AOF, there’s an unexpected behavior if stores key with expiration.

Let’s review by timeline:

  1. 00:00:00: SET foo bar EX 10, it stores into both AOF and storage.
  2. 00:00:05: GET foo, get bar because the key havn’t expired.
  3. 00:00:11: GET foo, get (nil) because the key have expired.
  4. 00:01:00: Shutdown service
  5. 01:00:00: Startup the service, read AOF and restore.
  6. 01:00:05: GET foo, get bar, it is an unexpected behavior.

To solve this problem, it should stores unix timestamp in AOF instead.

resp.Command{} with aof

We know it stores different command in AOF when SET with expiration. It needs to add aof in resp.Command{} because AOF writer uses this.

type Command struct {
	raw  Value
+	aof  Array
	cmd  BulkString
	args []BulkString
}

The aof is a deep clone of raw, but we’ve checked raw is resp.Array type in resp.ReadCommand().

func ReadCommand(rd *Reader) (*Command, error) {
	v, err := rd.Read()
	if err != nil {
		return nil, err
	}

	arr, ok := v.(Array)
	if !ok {
		return nil, fmt.Errorf("%w: expected array, got %T(%+v)", ErrProtocol, v, v)
	}
	if arr.null || len(arr.data) == 0 {
		return nil, fmt.Errorf("%w: empty array", ErrProtocol)
	}

	cmd, ok := arr.data[0].(BulkString)
	if !ok {
		return nil, fmt.Errorf("%w: command expected bulk string, got %T(%+v)", ErrProtocol, arr.data[0], arr.data[0])
	}

	args := make([]BulkString, 0, len(arr.data)-1)
	for i := range arr.data[1:] {
		arg, ok := arr.data[i+1].(BulkString)
		if !ok {
			return nil, fmt.Errorf("%w: argument [%d] expected bulk string, got %T(%+v)", ErrProtocol, i, arr.data[i+1], arr.data[i+1])
		}
		args = append(args, arg)
	}

	return &Command{
		raw:  v,
		aof:  arr.Clone(),
		cmd:  cmd,
		args: args,
	}, nil
}

Add add a UpdateAOF() to update aof:

func (cmd *Command) UpdateAOF(i int, v Value) {
	cmd.aof.data[i] = v
}

UpdateAOF() in SET command

	if p.Exp != nil && !p.Exp.IsZero() {
		cmd.UpdateAOF(expirationOptionIndex+1, resp.NewBulkString("PXAT"))
		cmd.UpdateAOF(expirationOptionIndex+2, resp.NewBulkString(strconv.FormatInt(p.Exp.UnixMilli(), 10)))
	}

It updates EX or PX argument into PXAT when storing AOF.

The TTL command

Implement and register a TTL Redis command:

package cmd

import (
	"context"
	"errors"
	"fmt"
	"time"

	"olivine/internal/repo"
	"olivine/pkg/resp"
)

type TTL struct {
	storage repo.Storage
}

func NewTTL(storage repo.Storage) *TTL {
	return &TTL{
		storage: storage,
	}
}

func (c *TTL) Command() string {
	return "TTL"
}

func (c *TTL) Exec(ctx context.Context, cmd *resp.Command) (resp.Value, error) {
	if err := c.parse(cmd); err != nil {
		return nil, err
	}

	args := cmd.Args()

	k := args[0]
	v, err := c.storage.Get(ctx, k.String())
	if err != nil {
		if errors.Is(err, repo.ErrNotFound) {
			return resp.Integer(-2), nil
		}
		return nil, fmt.Errorf("%w: %w", ErrStorage, err)
	}
	if v.ExpiresAt() == nil {
		return resp.Integer(-1), nil
	}

	return resp.Integer(time.Until(*v.ExpiresAt()) / time.Second), nil
}

func (c *TTL) parse(cmd *resp.Command) error {
	const (
		awaitingKey = iota
		keyReceived
		tooManyKeys
	)

	state := awaitingKey
	for range cmd.Args() {
		switch state {
		case awaitingKey:
			state = keyReceived
		case keyReceived:
			state = tooManyKeys
		}
	}

	if state != keyReceived {
		return fmt.Errorf("%w: argument count mismatch: expect %d got %d", ErrSyntax, 1, len(cmd.Args()))
	}

	return nil
}
// file: internal/service/cmd/cmd.go
func NewCommands(storage repo.Storage) []Command {
	return []Command{
		&Ping{},

		NewSet(storage),
		NewGet(storage),

+		NewTTL(storage),
	}
}

Note

This topic is hard to write. The key expiration feature is based on lots of computer science knowledge, but it could make this topic too long to read.

I want to explian why Redis uses both active and passive way to expire keys, and why not to use min-heap (priority queue) to handle a sorted array by expiration *time.Time for expiration.

I also want to explain how Redis stores expiration in redisObject by bitwise field likes magic, and I hope I can tell you why parse redis command syntax by self-made finite state machine in this topic, and why I give up to use regular expression.

I committed many changes not only in this topic, but I’m sorry that I cannot tell you more about these changes.