cachefs/pkg/fs/preload.go

// Copyright (C) 2022 Marius Schellenberger

package fs

import (
	"context"
	"io"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-logr/logr"
	"golang.org/x/exp/slices"
)

type queue []*Preload

func (queue) LessReverse(i, j *Preload) bool { return i.Prio > j.Prio }

func (q *queue) Sort() {
	slices.SortFunc(*q, q.LessReverse)
}

func (q *queue) Add(ph *PreloadHandler, name string, prio int) {
	for _, p := range *q {
		if p.Name == name {
			p.Prio += prio
			if p.Prio < 0 {
				p.Prio = 0
			}
			q.Sort()
			return
		}
	}
	*q = append(*q, &Preload{Name: name, ph: ph})
	q.Sort()
}

func (q *queue) Remove(name string, collect bool) {
	for i, p := range *q {
		if p.Name == name {
			p.stop(collect)
			*q = slices.Delete(*q, i, i+1)
			q.Sort()
			return
		}
	}
}

func (q *queue) GetPreload(name string) *Preload {
	for _, p := range *q {
		if p.Name == name {
			return p
		}
	}
	return nil
}

type Rate struct {
	r        io.Reader
	done     chan struct{}
	rate     atomic.Pointer[string]
	num      int64
	last     int64
	interval int
}

var zeroRate = fmtRate(0)

func newRate() *Rate {
	r := &Rate{
		done:     make(chan struct{}),
		interval: 10,
	}
	r.rate.Store(&zeroRate)
	return r
}

func (r *Rate) SetReader(reader io.Reader) {
	r.r = reader
	go r.sampler()
}

func (r *Rate) Read(b []byte) (n int, err error) {
	n, err = r.r.Read(b)
	atomic.AddInt64(&r.num, int64(n))
	return
}

func fmtRate(r int) string {
	switch {
	case r >= GiB:
		return strconv.FormatFloat(float64(r)/GiB, 'f', 2, 64) + " GiB/s"
	case r >= MiB:
		return strconv.FormatFloat(float64(r)/MiB, 'f', 2, 64) + " MiB/s"
	case r >= KiB:
		return strconv.FormatFloat(float64(r)/KiB, 'f', 2, 64) + " KiB/s"
	}
	return strconv.Itoa(r) + " B/s"
}

func (r *Rate) sampler() {
	t := time.NewTicker(time.Duration(r.interval) * time.Second)
	for {
		select {
		case <-r.done:
			t.Stop()
			return
		case <-t.C:
			num := atomic.LoadInt64(&r.num)
			rate := int(num-r.last) / r.interval
			r.last = num
			s := fmtRate(rate)
			r.rate.Store(&s)
		}
	}
}

func (r *Rate) String() string {
	s := r.rate.Load()
	if s == nil {
		return ""
	}
	return *s
}

func (r *Rate) Stop() {
	close(r.done)
	r.rate.Store(&zeroRate)
}

type Preload struct {
	ph      *PreloadHandler
	rate    *Rate
	delay   time.Duration
	Name    string
	cancel  func()
	Prio    int
	Status  int
	Errc    int
	Size    int64
	Running bool
}

type Preloads []Preload

type PreloadHandler struct {
	mu   sync.RWMutex
	log  logr.Logger
	fs   *FS
	q    queue
	fin  chan string
	err  chan string
	stop chan struct{}
	done chan struct{}
	max  int
}

func NewPreloadHandler(ctx context.Context, fs *FS, max int, log logr.Logger) (ph *PreloadHandler) {
	if max <= 0 {
		max = 1
	}
	ph = &PreloadHandler{
		log:  log,
		fs:   fs,
		q:    make(queue, 0),
		fin:  make(chan string, 1),
		err:  make(chan string, 1),
		stop: make(chan struct{}, 1),
		done: make(chan struct{}),
		max:  max,
	}
	go ph.preloadStatus(ctx)
	return
}

func (ph *PreloadHandler) RemovePreload(name string) {
	ph.mu.RLock()
	defer ph.mu.RUnlock()
	ph.q.Remove(name, true)
	ph.schedule()
}

func (ph *PreloadHandler) Preloads() (ps Preloads) {
	ph.mu.RLock()
	defer ph.mu.RUnlock()
	ps = make(Preloads, len(ph.q))
	for i, v := range ph.q {
		ps[i] = *v
		ps[i].Status = ph.fs.CacheStatus(v.Name)
	}
	return
}

func (ph *PreloadHandler) Preload(name string, prio int) {
	ph.mu.Lock()
	defer ph.mu.Unlock()
	ph.q.Add(ph, name, prio)
	ph.schedule()
}

func (ph *PreloadHandler) schedule() {
	for i := 0; i < len(ph.q); i++ {
		if i < ph.max {
			go ph.q[i].start()
		} else {
			ph.q[i].stop(true)
		}
	}
}

func (p *Preload) Rate() string {
	if p.rate != nil {
		return p.rate.String()
	}
	return ""
}

func (p *Preload) start() {
	if p.Running {
		return
	}
	if p.delay != 0 {
		time.Sleep(p.delay)
	}
	name := p.Name
	file, err := p.ph.fs.Open(name)
	if err != nil {
		p.ph.log.Error(err, "error staring next preload", "file", name)
		p.ph.err <- name
		return
	}
	f, ok := file.(*File)
	if !ok {
		return
	}
	ctx, cancel := context.WithCancel(context.Background())
	p.Running = true
	p.cancel = cancel
	p.rate = newRate()
	p.Size = f.size()
	go f.Preload(ctx,
		p.rate,
		func() {
			p.ph.fin <- name
		},
		func() {
			p.ph.err <- name
		},
		func() {
			p.ph.stop <- struct{}{}
		},
	)
}

func (p *Preload) stop(collect bool) {
	if p.Running {
		p.Running = false
		p.cancel()
		if collect {
			<-p.ph.stop
		}
	}
}

func (ph *PreloadHandler) preloadStatus(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			ph.close()
			return
		case name := <-ph.fin:
			ph.mu.Lock()
			ph.q.Remove(name, false)
			ph.schedule()
			ph.mu.Unlock()
		case name := <-ph.err:
			ph.mu.Lock()
			p := ph.q.GetPreload(name)
			if p != nil {
				p.stop(false)
				p.delay = time.Second * 2
				if p.Errc >= 10 {
					if p.Prio >= 0 {
						p.Prio = 0
					}
					p.Prio -= 1
					ph.schedule()
				} else {
					p.Errc += 1
					go p.start()
				}
			}
			ph.mu.Unlock()
		}
	}
}

func (ph *PreloadHandler) close() {
	ph.mu.Lock()
	defer ph.mu.Unlock()
	n := len(ph.q)
	for _, p := range ph.q {
		p.stop(false)
	}
	for i := 0; i < n; i++ {
		select {
		case <-ph.fin:
		case <-ph.err:
		case <-ph.stop:
		}
	}
	close(ph.done)
}

func (ph *PreloadHandler) Done() <-chan struct{} {
	return ph.done
}