Revive the rate-limited filesystem post-rebase

Author: bbockelm

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