From 6ff90fa0eef961d2e14fe8f23f33d553b11eb1d0 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 20 Jun 2022 11:28:45 +0100
Subject: [PATCH] Add `config.DataUnit` for specifying friendly cache sizes

---
 internal/caching/impl_ristretto.go |  7 ++++---
 setup/config/config_global.go      | 29 ++++++++++++++++++++++++++++-
 setup/config/config_test.go        | 21 +++++++++++++++++++++
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/internal/caching/impl_ristretto.go b/internal/caching/impl_ristretto.go
index 9284397df..3f63026ae 100644
--- a/internal/caching/impl_ristretto.go
+++ b/internal/caching/impl_ristretto.go
@@ -9,15 +9,16 @@ import (
 	"github.com/dgraph-io/ristretto"
 	"github.com/dgraph-io/ristretto/z"
 	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/dendrite/setup/config"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
-func MustCreateCache(maxCost int64, enablePrometheus bool) *ristretto.Cache {
+func MustCreateCache(maxCost config.DataUnit, enablePrometheus bool) *ristretto.Cache {
 	cache, err := ristretto.NewCache(&ristretto.Config{
 		NumCounters: 1e5,
-		MaxCost:     maxCost,
+		MaxCost:     int64(maxCost),
 		BufferItems: 64,
 		Metrics:     true,
 		KeyToHash: func(key interface{}) (uint64, uint64) {
@@ -56,7 +57,7 @@ const (
 	lazyLoadingCache
 )
 
-func NewRistrettoCache(maxCost int64, enablePrometheus bool) (*Caches, error) {
+func NewRistrettoCache(maxCost config.DataUnit, enablePrometheus bool) (*Caches, error) {
 	cache := MustCreateCache(maxCost, enablePrometheus)
 	return &Caches{
 		RoomVersions: &RistrettoCachePartition[string, gomatrixserverlib.RoomVersion]{
diff --git a/setup/config/config_global.go b/setup/config/config_global.go
index 1182b6953..343a23eaf 100644
--- a/setup/config/config_global.go
+++ b/setup/config/config_global.go
@@ -2,6 +2,8 @@ package config
 
 import (
 	"math/rand"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/matrix-org/gomatrixserverlib"
@@ -174,7 +176,7 @@ func (c *ServerNotices) Defaults(generate bool) {
 func (c *ServerNotices) Verify(errors *ConfigErrors, isMonolith bool) {}
 
 type Cache struct {
-	EstMaxSize int64 `yaml:"max_bytes_est"`
+	EstMaxSize DataUnit `yaml:"max_bytes_est"`
 }
 
 func (c *Cache) Defaults(generate bool) {
@@ -285,3 +287,28 @@ type PresenceOptions struct {
 	// Whether outbound presence events are allowed
 	EnableOutbound bool `yaml:"enable_outbound"`
 }
+
+type DataUnit int64
+
+func (d *DataUnit) UnmarshalText(text []byte) error {
+	var magnitude float64
+	s := strings.ToLower(string(text))
+	switch {
+	case strings.HasSuffix(s, "tb"):
+		s, magnitude = s[:len(s)-2], 1024*1024*1024*1024
+	case strings.HasSuffix(s, "gb"):
+		s, magnitude = s[:len(s)-2], 1024*1024*1024
+	case strings.HasSuffix(s, "mb"):
+		s, magnitude = s[:len(s)-2], 1024*1024
+	case strings.HasSuffix(s, "kb"):
+		s, magnitude = s[:len(s)-2], 1024
+	default:
+		magnitude = 1
+	}
+	v, err := strconv.ParseFloat(s, 64)
+	if err != nil {
+		return err
+	}
+	*d = DataUnit(v * magnitude)
+	return nil
+}
diff --git a/setup/config/config_test.go b/setup/config/config_test.go
index cbc57ad18..b9b1e7bb5 100644
--- a/setup/config/config_test.go
+++ b/setup/config/config_test.go
@@ -17,6 +17,8 @@ package config
 import (
 	"fmt"
 	"testing"
+
+	"gopkg.in/yaml.v2"
 )
 
 func TestLoadConfigRelative(t *testing.T) {
@@ -268,3 +270,22 @@ n0Xq64k7fc42HXJpF8CGBkSaIhtlzcruO+vqR80B9r62+D0V7VmHOnP135MT6noU
 ANAf5kxmMsM0zlN2hkxl0H6o7wKlBSw3RI3cjfilXiMWRPJrzlc4
 -----END CERTIFICATE-----
 `
+
+func TestUnmarshalDataUnit(t *testing.T) {
+	target := struct {
+		Got DataUnit `yaml:"value"`
+	}{}
+	for input, expect := range map[string]DataUnit{
+		"value: 0.6tb": 659706976665,
+		"value: 1.2gb": 1288490188,
+		"value: 256mb": 268435456,
+		"value: 128kb": 131072,
+		"value: 128":   128,
+	} {
+		if err := yaml.Unmarshal([]byte(input), &target); err != nil {
+			t.Fatal(err)
+		} else if target.Got != expect {
+			t.Fatalf("expected value %d but got %d", expect, target.Got)
+		}
+	}
+}