Skip to content

Commit 20bd8ec

Browse files
committed
mcp: use golang.org/x/time/rate for LoggingHandler rate limiting
Replace manual time-based rate limiting with golang.org/x/time/rate.Limiter for LoggingHandler. This provides a proper token bucket algorithm with better accuracy and standard library support. The rate limiter allows one log message per second with a burst of 5, matching the original intent while being more robust. Signed-off-by: wucm667 <stevenwucongmin@gmail.com> Closes #924
1 parent f5f2015 commit 20bd8ec

3 files changed

Lines changed: 12 additions & 12 deletions

File tree

go.mod

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ require (
99
github.com/segmentio/encoding v0.5.4
1010
github.com/yosida95/uritemplate/v3 v3.0.2
1111
golang.org/x/oauth2 v0.35.0
12+
golang.org/x/time v0.15.0
1213
golang.org/x/tools v0.42.0
1314
)
1415

go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,7 @@ golang.org/x/oauth2 v0.35.0 h1:Mv2mzuHuZuY2+bkyWXIHMfhNdJAdwW3FuWeCPYN5GVQ=
1414
golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
1515
golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
1616
golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
17+
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
18+
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
1719
golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
1820
golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=

mcp/logging.go

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ import (
1313
"slices"
1414
"sync"
1515
"time"
16+
17+
"golang.org/x/time/rate"
1618
)
1719

1820
// Logging levels.
@@ -83,9 +85,9 @@ type LoggingHandler struct {
8385
// Ensures that the buffer reset is atomic with the write (see Handle).
8486
// A pointer so that clones share the mutex. See
8587
// https://github.com/golang/example/blob/master/slog-handler-guide/README.md#getting-the-mutex-right.
86-
mu *sync.Mutex
87-
lastMessageSent time.Time // for rate-limiting
88-
buf *bytes.Buffer
88+
mu *sync.Mutex
89+
limiter *rate.Limiter // for rate-limiting
90+
buf *bytes.Buffer
8991
handler slog.Handler
9092
}
9193

@@ -118,6 +120,9 @@ func NewLoggingHandler(ss *ServerSession, opts *LoggingHandlerOptions) *LoggingH
118120
}
119121
if opts != nil {
120122
lh.opts = *opts
123+
if opts.MinInterval > 0 {
124+
lh.limiter = rate.NewLimiter(rate.Limit(1.0/opts.MinInterval.Seconds()), 1)
125+
}
121126
}
122127
return lh
123128
}
@@ -157,11 +162,7 @@ func (h *LoggingHandler) Handle(ctx context.Context, r slog.Record) error {
157162

158163
func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error {
159164
// Observe the rate limit.
160-
// TODO(jba): use golang.org/x/time/rate.
161-
h.mu.Lock()
162-
skip := time.Since(h.lastMessageSent) < h.opts.MinInterval
163-
h.mu.Unlock()
164-
if skip {
165+
if h.limiter != nil && !h.limiter.Allow() {
165166
return nil
166167
}
167168

@@ -184,10 +185,6 @@ func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error {
184185
return err
185186
}
186187

187-
h.mu.Lock()
188-
h.lastMessageSent = time.Now()
189-
h.mu.Unlock()
190-
191188
params := &LoggingMessageParams{
192189
Logger: h.opts.LoggerName,
193190
Level: slogLevelToMCP(r.Level),

0 commit comments

Comments
 (0)