Revive the rate-limited filesystem post-rebase

Author: bbockelm

generate/param_generator.go

@@ -83,6 +83,7 @@ func GenParamEnum() {
 	stringParamMap := make(map[string]string)
 	stringSliceParamMap := make(map[string]string)
 	intParamMap := make(map[string]string)
+	byteRateParamMap := make(map[string]string)
 	boolParamMap := make(map[string]string)
 	durationParamMap := make(map[string]string)
 	objectParamMap := make(map[string]string)
@@ -177,8 +178,7 @@ func GenParamEnum() {
 		case "int":
 			intParamMap[name] = rawName
 		case "byterate":
-			// byterate is stored as int (bytes per second) internally
-			intParamMap[name] = rawName
+			byteRateParamMap[name] = rawName
 		case "bool":
 			boolParamMap[name] = rawName
 		case "duration":
@@ -207,13 +207,14 @@ func GenParamEnum() {
 		StringMap              map[string]string
 		StringSliceMap         map[string]string
 		IntMap                 map[string]string
+		ByteRateMap            map[string]string
 		BoolMap                map[string]string
 		DurationMap            map[string]string
 		ObjectMap              map[string]string
 		DeprecatedMap          map[string][]string
 		RuntimeConfigurableMap map[string]bool
 		AllParamNames          []string
-	}{StringMap: stringParamMap, StringSliceMap: stringSliceParamMap, IntMap: intParamMap, BoolMap: boolParamMap, DurationMap: durationParamMap, ObjectMap: objectParamMap, DeprecatedMap: deprecatedMap, RuntimeConfigurableMap: runtimeConfigurableMap, AllParamNames: allParamNames})
+	}{StringMap: stringParamMap, StringSliceMap: stringSliceParamMap, IntMap: intParamMap, ByteRateMap: byteRateParamMap, BoolMap: boolParamMap, DurationMap: durationParamMap, ObjectMap: objectParamMap, DeprecatedMap: deprecatedMap, RuntimeConfigurableMap: runtimeConfigurableMap, AllParamNames: allParamNames})
 
 	if err != nil {
 		panic(err)
@@ -370,8 +371,7 @@ func GenParamStruct() {
 		case "int":
 			goType = "int"
 		case "byterate":
-			// byterate is parsed as string but stored as int (bytes per second)
-			goType = "int"
+			goType = "byte_rate.ByteRate"
 		case "bool":
 			goType = "bool"
 		case "duration":
@@ -459,6 +459,8 @@ import (
 	"time"
 
 	"github.com/spf13/viper"
+
+	"github.com/pelicanplatform/pelican/byte_rate"
 )
 
 type StringParam struct {
@@ -477,6 +479,10 @@ type IntParam struct {
 	name string
 }
 
+type ByteRateParam struct {
+	name string
+}
+
 type DurationParam struct {
 	name string
 }
@@ -604,6 +610,31 @@ func (iP IntParam) GetEnvVarName() string {
 	return paramNameToEnvVar(iP.name)
 }
 
+func (bRP ByteRateParam) GetByteRate() byte_rate.ByteRate {
+	config := getOrCreateConfig()
+	switch bRP.name {
+		case "Origin.TransferRateLimit":
+			return config.Origin.TransferRateLimit
+	}
+	return 0
+}
+
+func (bRP ByteRateParam) GetName() string {
+	return bRP.name
+}
+
+func (bRP ByteRateParam) IsSet() bool {
+	return viper.IsSet(bRP.name)
+}
+
+func (bRP ByteRateParam) IsRuntimeConfigurable() bool {
+	return IsRuntimeConfigurable(bRP.name)
+}
+
+func (bRP ByteRateParam) GetEnvVarName() string {
+	return paramNameToEnvVar(bRP.name)
+}
+
 func (bP BoolParam) GetBool() bool {
 	config := getOrCreateConfig()
 	switch bP.name {
@@ -702,6 +733,11 @@ var ({{range $key, $value := .IntMap}}
 	{{- end}}
 )
 
+var ({{range $key, $value := .ByteRateMap}}
+	{{$key}} = ByteRateParam{{"{"}}{{printf "%q" $value}}{{"}"}}
+	{{- end}}
+)
+
 var ({{range $key, $value := .BoolMap}}
 	{{$key}} = BoolParam{{"{"}}{{printf "%q" $value}}{{"}"}}
 	{{- end}}
@@ -741,6 +777,8 @@ package param
 
 import (
 	"time"
+
+	"github.com/pelicanplatform/pelican/byte_rate"
 )
 
 {{.GeneratedConfig}}

origin_serve/handlers.go

@@ -319,9 +319,9 @@ func InitializeHandlers(exports []server_utils.OriginExport) error {
 	exportPrefixMap = make(map[string]string) // Initialize the global map
 
 	// Get optional rate limit for testing
-	readRateLimit := param.Origin_TransferRateLimit.GetInt()
+	readRateLimit := param.Origin_TransferRateLimit.GetByteRate()
 	if readRateLimit > 0 {
-		log.Infof("Applying read rate limit: %d bytes/sec", readRateLimit)
+		log.Infof("Applying read rate limit: %s", readRateLimit.String())
 	}
 
 	for _, export := range exports {

origin_serve/rate_limited_fs.go

@@ -0,0 +1,123 @@
+/***************************************************************
+ *
+ * Copyright (C) 2026, Pelican Project, Morgridge Institute for Research
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you
+ * may not use this file except in compliance with the License.  You may
+ * obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ ***************************************************************/
+
+package origin_serve
+
+import (
+	"context"
+	"os"
+
+	"github.com/spf13/afero"
+	"golang.org/x/time/rate"
+
+	"github.com/pelicanplatform/pelican/byte_rate"
+)
+
+type rateLimitedFs struct {
+	afero.Fs
+	limiter *rate.Limiter
+}
+
+type rateLimitedFile struct {
+	afero.File
+	limiter *rate.Limiter
+}
+
+func newRateLimitedFs(fs afero.Fs, rateLimit byte_rate.ByteRate) afero.Fs {
+	if rateLimit <= 0 {
+		return fs
+	}
+	limit := int(rateLimit)
+	limiter := rate.NewLimiter(rate.Limit(limit), limit)
+	return &rateLimitedFs{Fs: fs, limiter: limiter}
+}
+
+func (r *rateLimitedFs) Open(name string) (afero.File, error) {
+	file, err := r.Fs.Open(name)
+	if err != nil {
+		return nil, err
+	}
+	return &rateLimitedFile{File: file, limiter: r.limiter}, nil
+}
+
+func (r *rateLimitedFs) OpenFile(name string, flag int, perm os.FileMode) (afero.File, error) {
+	file, err := r.Fs.OpenFile(name, flag, perm)
+	if err != nil {
+		return nil, err
+	}
+	return &rateLimitedFile{File: file, limiter: r.limiter}, nil
+}
+
+func (f *rateLimitedFile) Read(p []byte) (int, error) {
+	n, err := f.File.Read(p)
+	if n > 0 {
+		if waitErr := f.waitN(n); waitErr != nil {
+			return n, waitErr
+		}
+	}
+	return n, err
+}
+
+func (f *rateLimitedFile) Write(p []byte) (int, error) {
+	n, err := f.File.Write(p)
+	if n > 0 {
+		if waitErr := f.waitN(n); waitErr != nil {
+			return n, waitErr
+		}
+	}
+	return n, err
+}
+
+func (f *rateLimitedFile) ReadAt(p []byte, off int64) (int, error) {
+	n, err := f.File.ReadAt(p, off)
+	if n > 0 {
+		if waitErr := f.waitN(n); waitErr != nil {
+			return n, waitErr
+		}
+	}
+	return n, err
+}
+
+func (f *rateLimitedFile) WriteAt(p []byte, off int64) (int, error) {
+	n, err := f.File.WriteAt(p, off)
+	if n > 0 {
+		if waitErr := f.waitN(n); waitErr != nil {
+			return n, waitErr
+		}
+	}
+	return n, err
+}
+
+func (f *rateLimitedFile) waitN(n int) error {
+	remaining := n
+	for remaining > 0 {
+		burst := f.limiter.Burst()
+		if burst <= 0 {
+			burst = 1
+		}
+		step := remaining
+		if step > burst {
+			step = burst
+		}
+		if err := f.limiter.WaitN(context.Background(), step); err != nil {
+			return err
+		}
+		remaining -= step
+	}
+	return nil
+}

param/param.go

@@ -123,7 +123,7 @@ func stringToSliceHookFunc() mapstructure.DecodeHookFunc {
 	}
 }
 
-// stringToByteRateHookFunc returns a DecodeHookFunc that converts strings to integers
+// stringToByteRateHookFunc returns a DecodeHookFunc that converts strings to byte rates
 // representing bytes per second. It supports human-readable rate formats like:
 //   - "10MB/s", "100Mbps", "1.5GiB/m"
 //
@@ -132,8 +132,9 @@ func stringToSliceHookFunc() mapstructure.DecodeHookFunc {
 // so other hooks or default conversions can handle them.
 func stringToByteRateHookFunc() mapstructure.DecodeHookFunc {
 	return func(f reflect.Type, t reflect.Type, data interface{}) (interface{}, error) {
-		// Only convert string to int
-		if f.Kind() != reflect.String || t.Kind() != reflect.Int {
+		// Only convert string to int or ByteRate
+		byteRateType := reflect.TypeOf(byte_rate.ByteRate(0))
+		if f.Kind() != reflect.String || (t.Kind() != reflect.Int && t != byteRateType) {
 			return data, nil
 		}
 
@@ -144,6 +145,9 @@ func stringToByteRateHookFunc() mapstructure.DecodeHookFunc {
 
 		// Empty or "0" means no rate limiting
 		if raw == "" || raw == "0" {
+			if t == byteRateType {
+				return byte_rate.ByteRate(0), nil
+			}
 			return 0, nil
 		}
 
@@ -169,7 +173,10 @@ func stringToByteRateHookFunc() mapstructure.DecodeHookFunc {
 			return nil, errors.Wrapf(err, "failed to parse byte rate '%s'", raw)
 		}
 
-		// Return as integer bytes per second
+		if t == byteRateType {
+			return rate, nil
+		}
+		// Return as integer bytes per second for int targets
 		return int(rate), nil
 	}
 }

param/param_test.go

@@ -29,6 +29,8 @@ import (
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/pelicanplatform/pelican/byte_rate"
 )
 
 func TestSetAndGet(t *testing.T) {
@@ -663,7 +665,7 @@ func TestByteRateDecoding(t *testing.T) {
 		require.NoError(t, err)
 
 		// 10MB/s should be 10 * 1048576 (MiB) = 10485760 bytes/second
-		expected := 10 * 1048576
+		expected := byte_rate.ByteRate(10 * 1048576)
 		assert.Equal(t, expected, cfg.Origin.TransferRateLimit, "Should decode 10MB/s correctly")
 	})
 
@@ -688,7 +690,7 @@ func TestByteRateDecoding(t *testing.T) {
 		require.NoError(t, err)
 
 		// 100Mbps = 100 * 1048576 / 8 = 13107200 bytes/second
-		expected := 100 * 1048576 / 8
+		expected := byte_rate.ByteRate(100 * 1048576 / 8)
 		assert.Equal(t, expected, cfg.Origin.TransferRateLimit, "Should decode 100Mbps correctly")
 	})
 
@@ -712,7 +714,7 @@ func TestByteRateDecoding(t *testing.T) {
 		cfg, err := DecodeConfig(v)
 		require.NoError(t, err)
 
-		assert.Equal(t, 0, cfg.Origin.TransferRateLimit, "Should handle zero rate")
+		assert.Equal(t, byte_rate.ByteRate(0), cfg.Origin.TransferRateLimit, "Should handle zero rate")
 	})
 
 	t.Run("decode-invalid-rate-should-error", func(t *testing.T) {
@@ -747,7 +749,7 @@ func TestByteRateDecoding(t *testing.T) {
 		require.NoError(t, err)
 
 		// 5GB/s = 5 * 1073741824 = 5368709120 bytes/second
-		expected := 5 * 1073741824
+		expected := byte_rate.ByteRate(5 * 1073741824)
 		assert.Equal(t, expected, cfg.Origin.TransferRateLimit, "Should decode rate set via viper.Set")
 	})
 
@@ -763,8 +765,8 @@ func TestByteRateDecoding(t *testing.T) {
 
 		// Accessor should work even without explicit config creation
 		// because getOrCreateConfig will create it
-		rateLimit := Origin_TransferRateLimit.GetInt()
-		expected := 50 * 1048576
+		rateLimit := Origin_TransferRateLimit.GetByteRate()
+		expected := byte_rate.ByteRate(50 * 1048576)
 		assert.Equal(t, expected, rateLimit, "Accessor should return correct byte rate value")
 
 		// Verify config was created and stored