geo/server/internal/shared/cache/readthrough.go

package cache

import (
	"context"
	"encoding/json"
	"errors"
	"math/rand"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"
)

const (
	readthroughDefaultJitterRatio = 0.1
)

var (
	jitterMu   sync.Mutex
	jitterRand = rand.New(rand.NewSource(time.Now().UnixNano()))
)

type payloadEnvelope[T any] struct {
	Empty bool `json:"empty,omitempty"`
	Value T    `json:"value,omitempty"`
}

func LoadJSON[T any](
	ctx context.Context,
	c Cache,
	group *singleflight.Group,
	key string,
	ttl time.Duration,
	loader func(context.Context) (T, error),
) (T, error) {
	value, _, err := loadJSONInternal(ctx, c, group, key, ttl, 0, func(loadCtx context.Context) (T, bool, error) {
		value, err := loader(loadCtx)
		return value, true, err
	})
	return value, err
}

func LoadJSONWithEmpty[T any](
	ctx context.Context,
	c Cache,
	group *singleflight.Group,
	key string,
	ttl time.Duration,
	emptyTTL time.Duration,
	loader func(context.Context) (T, bool, error),
) (T, bool, error) {
	return loadJSONInternal(ctx, c, group, key, ttl, emptyTTL, loader)
}

func loadJSONInternal[T any](
	ctx context.Context,
	c Cache,
	group *singleflight.Group,
	key string,
	ttl time.Duration,
	emptyTTL time.Duration,
	loader func(context.Context) (T, bool, error),
) (T, bool, error) {
	var zero T
	if c == nil {
		recordCacheOperation("readthrough_load", "cache_disabled", time.Duration(0))
		return loader(ctx)
	}

	if raw, err := c.Get(ctx, key); err == nil {
		var envelope payloadEnvelope[T]
		if unmarshalErr := json.Unmarshal(raw, &envelope); unmarshalErr == nil {
			if envelope.Empty {
				recordCacheOperation("readthrough_empty", "hit", time.Duration(0))
				return zero, false, nil
			}
			recordCacheOperation("readthrough_hit", "ok", time.Duration(0))
			return envelope.Value, true, nil
		}
		recordCacheOperation("readthrough_unmarshal", "error", time.Duration(0))
	} else if errors.Is(err, ErrNotFound) {
		recordCacheOperation("readthrough_miss", "not_found", time.Duration(0))
	} else {
		recordCacheOperation("readthrough_get", "error", time.Duration(0))
	}

	result, err, _ := group.Do(key, func() (interface{}, error) {
		start := time.Now()
		value, found, loadErr := loader(ctx)
		recordCacheOperation("readthrough_load", cacheResult(loadErr), time.Since(start))
		if loadErr != nil {
			return nil, loadErr
		}

		envelope := payloadEnvelope[T]{
			Empty: !found,
			Value: value,
		}
		if raw, marshalErr := json.Marshal(envelope); marshalErr == nil {
			recordCacheOperation("readthrough_marshal", "ok", time.Duration(0))
			cacheTTL := ttl
			if !found {
				cacheTTL = emptyTTL
			}
			if cacheTTL > 0 {
				setStart := time.Now()
				setErr := c.Set(ctx, key, raw, ApplyJitterWithRatio(cacheTTL, cacheJitterRatio(c)))
				recordCacheOperation("readthrough_set", cacheResult(setErr), time.Since(setStart))
			}
		} else {
			recordCacheOperation("readthrough_marshal", "error", time.Duration(0))
		}

		return envelope, nil
	})
	if err != nil {
		return zero, false, err
	}

	envelope, ok := result.(payloadEnvelope[T])
	if !ok {
		return zero, false, nil
	}
	if envelope.Empty {
		return zero, false, nil
	}
	return envelope.Value, true, nil
}

func ApplyJitter(ttl time.Duration) time.Duration {
	return ApplyJitterWithRatio(ttl, readthroughDefaultJitterRatio)
}

func ApplyJitterWithRatio(ttl time.Duration, ratio float64) time.Duration {
	if ttl <= 0 {
		return ttl
	}
	if ratio <= 0 {
		return ttl
	}
	if ratio > 1 {
		ratio = 1
	}

	maxJitter := int64(float64(ttl) * ratio)
	if maxJitter <= 0 {
		return ttl
	}

	jitterMu.Lock()
	defer jitterMu.Unlock()

	return ttl + time.Duration(jitterRand.Int63n(maxJitter+1))
}

func cacheJitterRatio(c Cache) float64 {
	if c == nil {
		return readthroughDefaultJitterRatio
	}
	if provider, ok := c.(optionsProvider); ok {
		return provider.CacheOptions().JitterRatio
	}
	return readthroughDefaultJitterRatio
}