@@ -19,6 +19,7 @@ import (
1919 "github.com/sei-protocol/sei-chain/sei-db/common/logger"
2020 "github.com/sei-protocol/sei-chain/sei-db/state_db/sc/types"
2121 "golang.org/x/sys/unix"
22+ "golang.org/x/time/rate"
2223)
2324
2425const (
@@ -53,6 +54,66 @@ func (w *monitoringWriter) Write(p []byte) (n int, err error) {
5354 return n , err
5455}
5556
57+ // rateLimitedWriter wraps an io.Writer with rate limiting to prevent
58+ // page cache eviction on machines with limited RAM.
59+ type rateLimitedWriter struct {
60+ w io.Writer
61+ limiter * rate.Limiter
62+ ctx context.Context
63+ }
64+
65+ // NewGlobalRateLimiter creates a shared rate limiter for snapshot writes.
66+ // rateMBps is the rate limit in MB/s. If <= 0, returns nil (no limit).
67+ // This limiter should be shared across all files and trees in a single snapshot operation.
68+ func NewGlobalRateLimiter (rateMBps int ) * rate.Limiter {
69+ if rateMBps <= 0 {
70+ return nil
71+ }
72+ const mb = 1024 * 1024
73+ bytesPerSec := rate .Limit (rateMBps * mb )
74+ // Burst = 4MB: small enough to spread large bufio flushes (128MB) across
75+ // many smaller IO ops, preventing page cache eviction spikes.
76+ burstBytes := 4 * mb
77+ return rate .NewLimiter (bytesPerSec , burstBytes )
78+ }
79+
80+ // newRateLimitedWriter creates a rate-limited writer with a shared limiter.
81+ // If limiter is nil, returns the original writer (no limit).
82+ func newRateLimitedWriter (ctx context.Context , w io.Writer , limiter * rate.Limiter ) io.Writer {
83+ if limiter == nil {
84+ return w
85+ }
86+ return & rateLimitedWriter {
87+ w : w ,
88+ limiter : limiter ,
89+ ctx : ctx ,
90+ }
91+ }
92+
93+ func (w * rateLimitedWriter ) Write (p []byte ) (n int , err error ) {
94+ // Wait for rate limiter before writing
95+ // For large writes, we may need to wait multiple times
96+ remaining := len (p )
97+ written := 0
98+ for remaining > 0 {
99+ // Limit each wait to burst size to avoid very long waits
100+ toWrite := remaining
101+ if toWrite > w .limiter .Burst () {
102+ toWrite = w .limiter .Burst ()
103+ }
104+ if err := w .limiter .WaitN (w .ctx , toWrite ); err != nil {
105+ return written , err
106+ }
107+ n , err := w .w .Write (p [written : written + toWrite ])
108+ written += n
109+ remaining -= n
110+ if err != nil {
111+ return written , err
112+ }
113+ }
114+ return written , nil
115+ }
116+
56117// Snapshot manage the lifecycle of mmap-ed files for the snapshot,
57118// it must out live the objects that derived from it.
58119type Snapshot struct {
@@ -390,6 +451,12 @@ func (snapshot *Snapshot) export(callback func(*types.SnapshotNode) bool) {
390451}
391452
392453func (t * Tree ) WriteSnapshot (ctx context.Context , snapshotDir string ) error {
454+ return t .WriteSnapshotWithRateLimit (ctx , snapshotDir , nil )
455+ }
456+
457+ // WriteSnapshotWithRateLimit writes snapshot with optional rate limiting.
458+ // limiter is a shared rate limiter. nil means unlimited.
459+ func (t * Tree ) WriteSnapshotWithRateLimit (ctx context.Context , snapshotDir string , limiter * rate.Limiter ) error {
393460 // Estimate tree size: root.Size() returns leaf count, total = leaves + branches ≈ 2x
394461 treeSize := int64 (0 )
395462 if t .root != nil {
@@ -399,7 +466,7 @@ func (t *Tree) WriteSnapshot(ctx context.Context, snapshotDir string) error {
399466 // Use 128MB buffer for all trees (large buffer for better performance)
400467 bufSize := bufIOSize
401468
402- err := writeSnapshotWithBuffer (ctx , snapshotDir , t .version , bufSize , treeSize , t .logger , func (w * snapshotWriter ) (uint32 , error ) {
469+ err := writeSnapshotWithBuffer (ctx , snapshotDir , t .version , bufSize , treeSize , limiter , t .logger , func (w * snapshotWriter ) (uint32 , error ) {
403470 if t .root == nil {
404471 return 0 , nil
405472 }
@@ -417,12 +484,14 @@ func (t *Tree) WriteSnapshot(ctx context.Context, snapshotDir string) error {
417484 return nil
418485}
419486
420- // writeSnapshotWithBuffer writes snapshot with specified buffer size
487+ // writeSnapshotWithBuffer writes snapshot with specified buffer size and optional rate limiting.
488+ // limiter is a shared rate limiter. nil means unlimited.
421489func writeSnapshotWithBuffer (
422490 ctx context.Context ,
423491 dir string , version uint32 ,
424492 bufSize int ,
425493 totalNodes int64 ,
494+ limiter * rate.Limiter ,
426495 log logger.Logger ,
427496 doWrite func (* snapshotWriter ) (uint32 , error ),
428497) (returnErr error ) {
@@ -469,10 +538,17 @@ func writeSnapshotWithBuffer(
469538 leavesMonitor := & monitoringWriter {f : fpLeaves }
470539 kvsMonitor := & monitoringWriter {f : fpKVs }
471540
541+ // Apply rate limiting if configured (shared limiter across all files)
542+ // This ensures total write rate is capped regardless of file count
543+ var nodesRateLimited , leavesRateLimited , kvsRateLimited io.Writer
544+ nodesRateLimited = newRateLimitedWriter (ctx , nodesMonitor , limiter )
545+ leavesRateLimited = newRateLimitedWriter (ctx , leavesMonitor , limiter )
546+ kvsRateLimited = newRateLimitedWriter (ctx , kvsMonitor , limiter )
547+
472548 // Create buffered writers with buffers
473- nodesWriter := bufio .NewWriterSize (nodesMonitor , bufSize )
474- leavesWriter := bufio .NewWriterSize (leavesMonitor , bufSize )
475- kvsWriter := bufio .NewWriterSize (kvsMonitor , bufSize )
549+ nodesWriter := bufio .NewWriterSize (nodesRateLimited , bufSize )
550+ leavesWriter := bufio .NewWriterSize (leavesRateLimited , bufSize )
551+ kvsWriter := bufio .NewWriterSize (kvsRateLimited , bufSize )
476552
477553 w := newSnapshotWriter (ctx , nodesWriter , leavesWriter , kvsWriter , log )
478554 w .treeName = filepath .Base (dir ) // Set tree name for progress reporting
@@ -546,8 +622,8 @@ func writeSnapshot(
546622 dir string , version uint32 ,
547623 doWrite func (* snapshotWriter ) (uint32 , error ),
548624) error {
549- // Use nop logger for backward compatibility
550- return writeSnapshotWithBuffer (ctx , dir , version , bufIOSize , 0 , logger .NewNopLogger (), doWrite )
625+ // Use nop logger and no rate limit for backward compatibility
626+ return writeSnapshotWithBuffer (ctx , dir , version , bufIOSize , 0 , nil , logger .NewNopLogger (), doWrite )
551627}
552628
553629// kvWriteOp represents a key-value write operation
0 commit comments