tile38/internal/server/aof.go

package server

import (
	"errors"
	"fmt"
	"io"
	"math"
	"net"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/tidwall/buntdb"
	"github.com/tidwall/gjson"
	"github.com/tidwall/redcon"
	"github.com/tidwall/resp"
	"github.com/tidwall/tile38/internal/log"
)

type errAOFHook struct {
	err error
}

func (err errAOFHook) Error() string {
	return fmt.Sprintf("hook: %v", err.err)
}

func (s *Server) loadAOF() (err error) {
	fi, err := s.aof.Stat()
	if err != nil {
		return err
	}
	start := time.Now()
	var count int
	defer func() {
		d := time.Since(start)
		ps := float64(count) / (float64(d) / float64(time.Second))
		suf := []string{"bytes/s", "KB/s", "MB/s", "GB/s", "TB/s"}
		bps := float64(fi.Size()) / (float64(d) / float64(time.Second))
		for i := 0; bps > 1024 && len(suf) > 1; i++ {
			bps /= 1024
			suf = suf[1:]
		}
		byteSpeed := fmt.Sprintf("%.0f %s", bps, suf[0])
		log.Infof("AOF loaded %d commands: %.2fs, %.0f/s, %s",
			count, float64(d)/float64(time.Second), ps, byteSpeed)
	}()
	var buf []byte
	var args [][]byte
	var packet [0xFFFF]byte
	for {
		n, err := s.aof.Read(packet[:])
		if err != nil {
			if err != io.EOF {
				return err
			}
			if len(buf) > 0 {
				// There was an incomplete command or other data at the end of
				// the AOF file. Attempt to recover the file by truncating the
				// file at the end position of the last complete command.
				log.Warnf("Truncating %d bytes due to an incomplete command\n",
					len(buf))
				s.aofsz -= len(buf)
				if err := s.aof.Truncate(int64(s.aofsz)); err != nil {
					return err
				}
				if _, err := s.aof.Seek(int64(s.aofsz), 0); err != nil {
					return err
				}
			}
			return nil
		}
		s.aofsz += n
		data := packet[:n]
		if len(buf) > 0 {
			data = append(buf, data...)
		}
		var complete bool
		for {
			if len(data) > 0 && data[0] == 0 {
				// Zeros found in AOF file (issue #230).
				// Just ignore it and move the next byte.
				data = data[1:]
				continue
			}
			complete, args, _, data, err = redcon.ReadNextCommand(data, args[:0])
			if err != nil {
				return err
			}
			if !complete {
				break
			}
			if len(args) > 0 {
				var msg Message
				msg.Args = msg.Args[:0]
				for _, arg := range args {
					msg.Args = append(msg.Args, string(arg))
				}
				if _, _, err := s.command(&msg, nil); err != nil {
					if commandErrIsFatal(err) {
						return err
					}
				}
				count++
			}
		}
		if len(data) > 0 {
			buf = append(buf[:0], data...)
		} else if len(buf) > 0 {
			buf = buf[:0]
		}
	}
}

func commandErrIsFatal(err error) bool {
	// FSET (and other writable commands) may return errors that we need
	// to ignore during the loading process. These errors may occur (though unlikely)
	// due to the aof rewrite operation.
	return !(err == errKeyNotFound || err == errIDNotFound)
}

// flushAOF flushes all aof buffer data to disk. Set sync to true to sync the
// fsync the file.
func (s *Server) flushAOF(sync bool) {
	if len(s.aofbuf) > 0 {
		_, err := s.aof.Write(s.aofbuf)
		if err != nil {
			panic(err)
		}
		if sync {
			if err := s.aof.Sync(); err != nil {
				panic(err)
			}
		}
		if cap(s.aofbuf) > 1024*1024*32 {
			s.aofbuf = make([]byte, 0, 1024*1024*32)
		} else {
			s.aofbuf = s.aofbuf[:0]
		}
	}
}

func (s *Server) writeAOF(args []string, d *commandDetails) error {
	if d != nil && !d.updated {
		// just ignore writes if the command did not update
		return nil
	}

	if s.shrinking {
		nargs := make([]string, len(args))
		copy(nargs, args)
		s.shrinklog = append(s.shrinklog, nargs)
	}

	if s.aof != nil {
		s.aofdirty.Store(true) // prewrite optimization flag
		n := len(s.aofbuf)
		s.aofbuf = redcon.AppendArray(s.aofbuf, len(args))
		for _, arg := range args {
			s.aofbuf = redcon.AppendBulkString(s.aofbuf, arg)
		}
		s.aofsz += len(s.aofbuf) - n
	}

	// notify aof live connections that we have new data
	s.fcond.L.Lock()
	s.fcond.Broadcast()
	s.fcond.L.Unlock()

	// process geofences
	if d != nil {
		// webhook geofences
		if s.config.followHost() == "" {
			// for leader only
			if d.parent {
				// queue children
				for _, d := range d.children {
					if err := s.queueHooks(d); err != nil {
						return err
					}
				}
			} else {
				// queue parent
				if err := s.queueHooks(d); err != nil {
					return err
				}
			}
		}

		// live geofences
		s.lcond.L.Lock()
		if len(s.lives) > 0 {
			if d.parent {
				// queue children
				s.lstack = append(s.lstack, d.children...)
			} else {
				// queue parent
				s.lstack = append(s.lstack, d)
			}
			s.lcond.Broadcast()
		}
		s.lcond.L.Unlock()
	}
	return nil
}

func (s *Server) getQueueCandidates(d *commandDetails) []*Hook {
	candidates := make(map[*Hook]bool)
	// add the hooks with "outside" detection
	s.hooksOut.Ascend(nil, func(v interface{}) bool {
		hook := v.(*Hook)
		if hook.Key == d.key {
			candidates[hook] = true
		}
		return true
	})
	// look for candidates that might "cross" geofences
	if d.old != nil && d.obj != nil && s.hookCross.Len() > 0 {
		r1, r2 := d.old.Rect(), d.obj.Rect()
		s.hookCross.Search(
			[2]float64{
				math.Min(r1.Min.X, r2.Min.X),
				math.Min(r1.Min.Y, r2.Min.Y),
			},
			[2]float64{
				math.Max(r1.Max.X, r2.Max.X),
				math.Max(r1.Max.Y, r2.Max.Y),
			},
			func(min, max [2]float64, value interface{}) bool {
				hook := value.(*Hook)
				if hook.Key == d.key {
					candidates[hook] = true
				}
				return true
			})
	}
	// look for candidates that overlap the old object
	if d.old != nil {
		r1 := d.old.Rect()
		s.hookTree.Search(
			[2]float64{r1.Min.X, r1.Min.Y},
			[2]float64{r1.Max.X, r1.Max.Y},
			func(min, max [2]float64, value interface{}) bool {
				hook := value.(*Hook)
				if hook.Key == d.key {
					candidates[hook] = true
				}
				return true
			})
	}
	// look for candidates that overlap the new object
	if d.obj != nil {
		r1 := d.obj.Rect()
		s.hookTree.Search(
			[2]float64{r1.Min.X, r1.Min.Y},
			[2]float64{r1.Max.X, r1.Max.Y},
			func(min, max [2]float64, value interface{}) bool {
				hook := value.(*Hook)
				if hook.Key == d.key {
					candidates[hook] = true
				}
				return true
			})
	}
	if len(candidates) == 0 {
		return nil
	}
	// return the candidates as a slice
	ret := make([]*Hook, 0, len(candidates))
	for hook := range candidates {
		ret = append(ret, hook)
	}
	return ret
}

func (s *Server) queueHooks(d *commandDetails) error {
	// Create the slices that will store all messages and hooks
	var cmsgs, wmsgs []string
	var whooks []*Hook

	// Compile a slice of potential hook recipients
	candidates := s.getQueueCandidates(d)
	for _, hook := range candidates {
		// Calculate all matching fence messages for all candidates and append
		// them to the appropriate message slice
		msgs := FenceMatch(hook.Name, hook.ScanWriter, hook.Fence, hook.Metas, d)
		if len(msgs) > 0 {
			if hook.channel {
				cmsgs = append(cmsgs, msgs...)
			} else {
				wmsgs = append(wmsgs, msgs...)
				whooks = append(whooks, hook)
			}
		}
	}

	// Return nil if there are no messages to be sent
	if len(cmsgs)+len(wmsgs) == 0 {
		return nil
	}

	// Sort both message channel and webhook message slices
	if len(cmsgs) > 1 {
		sortMsgs(cmsgs)
	}
	if len(wmsgs) > 1 {
		sortMsgs(wmsgs)
	}

	// Publish all channel messages if any exist
	if len(cmsgs) > 0 {
		for _, m := range cmsgs {
			s.Publish(gjson.Get(m, "hook").String(), m)
		}
	}

	// Queue the webhook messages in the buntdb database
	err := s.qdb.Update(func(tx *buntdb.Tx) error {
		for _, msg := range wmsgs {
			s.qidx++ // increment the log id
			key := hookLogPrefix + uint64ToString(s.qidx)
			_, _, err := tx.Set(key, msg, hookLogSetDefaults)
			if err != nil {
				return err
			}
			log.Debugf("queued hook: %d", s.qidx)
		}
		_, _, err := tx.Set("hook:idx", uint64ToString(s.qidx), nil)
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return err
	}
	// all the messages have been queued.
	// notify the hooks
	for _, hook := range whooks {
		hook.Signal()
	}
	return nil
}

// sortMsgs sorts passed notification messages by their detect and hook fields
func sortMsgs(msgs []string) {
	sort.SliceStable(msgs, func(i, j int) bool {
		detectI := msgDetectCode(gjson.Get(msgs[i], "detect").String())
		detectJ := msgDetectCode(gjson.Get(msgs[j], "detect").String())
		if detectI < detectJ {
			return true
		}
		if detectI > detectJ {
			return false
		}
		hookI := gjson.Get(msgs[i], "hook").String()
		hookJ := gjson.Get(msgs[j], "hook").String()
		return hookI < hookJ
	})
}

// msgDetectCode returns a weight value for the passed detect value
func msgDetectCode(detect string) int {
	switch detect {
	case "exit":
		return 1
	case "outside":
		return 2
	case "enter":
		return 3
	case "inside":
		return 4
	default:
		return 0
	}
}

// Converts string to an integer
func stringToUint64(s string) uint64 {
	n, _ := strconv.ParseUint(s, 10, 64)
	return n
}

// Converts a uint to a string
func uint64ToString(u uint64) string {
	s := strings.Repeat("0", 20) + strconv.FormatUint(u, 10)
	return s[len(s)-20:]
}

type liveAOFSwitches struct {
	pos int64
}

func (s liveAOFSwitches) Error() string {
	return goingLive
}

// AOFMD5 pos size
func (s *Server) cmdAOFMD5(msg *Message) (resp.Value, error) {
	start := time.Now()

	// >> Args

	args := msg.Args
	if len(args) != 3 {
		return retrerr(errInvalidNumberOfArguments)
	}
	pos, err := strconv.ParseInt(args[1], 10, 64)
	if err != nil || pos < 0 {
		return retrerr(errInvalidArgument(args[1]))
	}
	size, err := strconv.ParseInt(args[2], 10, 64)
	if err != nil || size < 0 {
		return retrerr(errInvalidArgument(args[2]))
	}

	// >> Operation

	sum, err := s.checksum(pos, size)
	if err != nil {
		return retrerr(err)
	}

	// >> Response

	if msg.OutputType == JSON {
		return resp.StringValue(fmt.Sprintf(
			`{"ok":true,"md5":"%s","elapsed":"%s"}`,
			sum, time.Since(start))), nil
	}
	return resp.SimpleStringValue(sum), nil
}

// AOF pos
func (s *Server) cmdAOF(msg *Message) (resp.Value, error) {
	if s.aof == nil {
		return retrerr(errors.New("aof disabled"))
	}

	// >> Args

	args := msg.Args
	if len(args) != 2 {
		return retrerr(errInvalidNumberOfArguments)
	}

	pos, err := strconv.ParseInt(args[1], 10, 64)
	if err != nil || pos < 0 {
		return retrerr(errInvalidArgument(args[1]))
	}

	// >> Operation

	f, err := os.Open(s.aof.Name())
	if err != nil {
		return retrerr(err)
	}
	defer f.Close()

	n, err := f.Seek(0, 2)
	if err != nil {
		return retrerr(err)
	}

	if n < pos {
		return retrerr(errors.New(
			"pos is too big, must be less that the aof_size of leader"))
	}

	// >> Response

	var ls liveAOFSwitches
	ls.pos = pos
	return NOMessage, ls
}

func (s *Server) liveAOF(pos int64, conn net.Conn, rd *PipelineReader, msg *Message) error {
	s.mu.RLock()
	f, err := os.Open(s.aof.Name())
	s.mu.RUnlock()
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.aofconnM[conn] = f
	s.mu.Unlock()
	defer func() {
		s.mu.Lock()
		delete(s.aofconnM, conn)
		s.mu.Unlock()
		conn.Close()
		f.Close()
	}()

	if _, err := conn.Write([]byte("+OK\r\n")); err != nil {
		return err
	}
	if _, err := f.Seek(pos, 0); err != nil {
		return err
	}
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer func() {
			f.Close()
			conn.Close()
			wg.Done()
		}()
		// Any incoming message should end the connection
		rd.ReadMessages()
	}()
	_, err = io.Copy(conn, f)
	if err != nil {
		return err
	}
	b := make([]byte, 4096*2)
	for {
		n, err := f.Read(b)
		if n > 0 {
			if _, err := conn.Write(b[:n]); err != nil {
				return err
			}
		}
		if err == io.EOF {
			time.Sleep(time.Second / 4)
		} else if err != nil {
			return err
		}
	}
}