mcp: use golang.org/x/time/rate for LoggingHandler rate limiting (#927)
Replaces the simple time-interval check in `LoggingHandler` with a standard `rate.Limiter` from `golang.org/x/time/rate`, as suggested by the existing TODO comment. This provides a well-tested implementation and removes the need for manual `lastMessageSent` tracking. The rate (`1/MinInterval`) and burst (1) match the previous behavior exactly. /cc @jba Closes #924 --------- Signed-off-by: wucm667 <stevenwucongmin@gmail.com> Co-authored-by: Jonathan Amsterdam <jba@users.noreply.github.com> Co-authored-by: Maciej Kisiel <mkisiel@google.com>
This commit is contained in:
parent
06c5f290f2
commit
9604499eed
3 changed files with 13 additions and 13 deletions
1
go.mod
1
go.mod
|
|
@ -9,6 +9,7 @@ require (
|
||||||
github.com/segmentio/encoding v0.5.4
|
github.com/segmentio/encoding v0.5.4
|
||||||
github.com/yosida95/uritemplate/v3 v3.0.2
|
github.com/yosida95/uritemplate/v3 v3.0.2
|
||||||
golang.org/x/oauth2 v0.35.0
|
golang.org/x/oauth2 v0.35.0
|
||||||
|
golang.org/x/time v0.15.0
|
||||||
golang.org/x/tools v0.42.0
|
golang.org/x/tools v0.42.0
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
2
go.sum
2
go.sum
|
|
@ -14,5 +14,7 @@ golang.org/x/oauth2 v0.35.0 h1:Mv2mzuHuZuY2+bkyWXIHMfhNdJAdwW3FuWeCPYN5GVQ=
|
||||||
golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
|
golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
|
||||||
golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
|
golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
|
||||||
golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||||
|
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
|
||||||
|
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
|
||||||
golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
|
golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
|
||||||
golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
|
golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,8 @@ import (
|
||||||
"slices"
|
"slices"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"golang.org/x/time/rate"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Logging levels.
|
// Logging levels.
|
||||||
|
|
@ -83,10 +85,10 @@ type LoggingHandler struct {
|
||||||
// Ensures that the buffer reset is atomic with the write (see Handle).
|
// Ensures that the buffer reset is atomic with the write (see Handle).
|
||||||
// A pointer so that clones share the mutex. See
|
// A pointer so that clones share the mutex. See
|
||||||
// https://github.com/golang/example/blob/master/slog-handler-guide/README.md#getting-the-mutex-right.
|
// https://github.com/golang/example/blob/master/slog-handler-guide/README.md#getting-the-mutex-right.
|
||||||
mu *sync.Mutex
|
mu *sync.Mutex
|
||||||
lastMessageSent time.Time // for rate-limiting
|
limiter *rate.Limiter // for rate-limiting
|
||||||
buf *bytes.Buffer
|
buf *bytes.Buffer
|
||||||
handler slog.Handler
|
handler slog.Handler
|
||||||
}
|
}
|
||||||
|
|
||||||
// ensureLogger returns l if non-nil, otherwise a discard logger.
|
// ensureLogger returns l if non-nil, otherwise a discard logger.
|
||||||
|
|
@ -118,6 +120,9 @@ func NewLoggingHandler(ss *ServerSession, opts *LoggingHandlerOptions) *LoggingH
|
||||||
}
|
}
|
||||||
if opts != nil {
|
if opts != nil {
|
||||||
lh.opts = *opts
|
lh.opts = *opts
|
||||||
|
if opts.MinInterval > 0 {
|
||||||
|
lh.limiter = rate.NewLimiter(rate.Every(opts.MinInterval), 1)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return lh
|
return lh
|
||||||
}
|
}
|
||||||
|
|
@ -157,11 +162,7 @@ func (h *LoggingHandler) Handle(ctx context.Context, r slog.Record) error {
|
||||||
|
|
||||||
func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error {
|
func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error {
|
||||||
// Observe the rate limit.
|
// Observe the rate limit.
|
||||||
// TODO(jba): use golang.org/x/time/rate.
|
if h.limiter != nil && !h.limiter.Allow() {
|
||||||
h.mu.Lock()
|
|
||||||
skip := time.Since(h.lastMessageSent) < h.opts.MinInterval
|
|
||||||
h.mu.Unlock()
|
|
||||||
if skip {
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -184,10 +185,6 @@ func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
h.mu.Lock()
|
|
||||||
h.lastMessageSent = time.Now()
|
|
||||||
h.mu.Unlock()
|
|
||||||
|
|
||||||
params := &LoggingMessageParams{
|
params := &LoggingMessageParams{
|
||||||
Logger: h.opts.LoggerName,
|
Logger: h.opts.LoggerName,
|
||||||
Level: slogLevelToMCP(r.Level),
|
Level: slogLevelToMCP(r.Level),
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue