rate-limited backend; context to propagate vlabs; prefetch

alex-aizman · alex-aizman · commit 8ee8b4425fc7 · 2025-03-14T15:58:03.000-04:00
* use context to propagate vlabs all the way to backend API - add setter and getter * prefetch: remove too-many retries - backend's responsibility going forward * part three, prev, commit: c4b796a Signed-off-by: Alex Aizman <alex.aizman@gmail.com>
diff --git a/ais/rlbackend.go b/ais/rlbackend.go
@@ -15,12 +15,11 @@ import (
 	"github.com/NVIDIA/aistore/core"
 	"github.com/NVIDIA/aistore/core/meta"
 	"github.com/NVIDIA/aistore/stats"
+	"github.com/NVIDIA/aistore/xact"
 )
 
 // NOTE: backend.GetObjReader is handled elsewhere - by xactions
 
-// TODO -- FIXME: nil context (below); callers to provide context
-
 type (
 	rlbackend struct {
 		core.Backend
@@ -29,60 +28,78 @@ type (
 )
 
 func (bp *rlbackend) GetObj(ctx context.Context, lom *core.LOM, owt cmn.OWT, origReq *http.Request) (ecode int, err error) {
+	ecode, err = bp.Backend.GetObj(ctx, lom, owt, origReq)
+	if err == nil {
+		return ecode, nil
+	}
+
 	var tot time.Duration
 	for tot < cos.DfltRateMaxWait {
-		ecode, err = bp.Backend.GetObj(ctx, lom, owt, origReq)
-		if err == nil {
-			return ecode, nil
-		}
-		sleep := bp.retry(ctx, err, lom.Bck(), http.MethodGet, stats.GetRateRetryLatencyTotal)
+		sleep := bp.retry(err, lom.Bck(), http.MethodGet)
 		if sleep == 0 {
-			return ecode, err
+			break
 		}
 		tot += sleep
+		ecode, err = bp.Backend.GetObj(ctx, lom, owt, origReq)
+		if err == nil {
+			break
+		}
 	}
+	bp.stats(ctx, lom.Bck(), stats.GetRateRetryLatencyTotal, tot)
 	return ecode, err
 }
 
 func (bp *rlbackend) PutObj(ctx context.Context, r io.ReadCloser, lom *core.LOM, origReq *http.Request) (ecode int, err error) {
+	ecode, err = bp.Backend.PutObj(ctx, r, lom, origReq)
+	if err == nil {
+		return ecode, nil
+	}
+
 	var tot time.Duration
 	for tot < cos.DfltRateMaxWait {
-		ecode, err = bp.Backend.PutObj(ctx, r, lom, origReq)
-		if err == nil {
-			return ecode, nil
-		}
-		sleep := bp.retry(ctx, err, lom.Bck(), http.MethodPut, stats.PutRateRetryLatencyTotal)
+		sleep := bp.retry(err, lom.Bck(), http.MethodPut)
 		if sleep == 0 {
-			return ecode, err
+			break
 		}
 		tot += sleep
+		ecode, err = bp.Backend.PutObj(ctx, r, lom, origReq)
+		if err == nil {
+			break
+		}
 	}
+	bp.stats(ctx, lom.Bck(), stats.PutRateRetryLatencyTotal, tot)
 	return ecode, err
 }
 
 func (bp *rlbackend) DeleteObj(ctx context.Context, lom *core.LOM) (ecode int, err error) {
+	ecode, err = bp.Backend.DeleteObj(ctx, lom)
+	if err == nil {
+		return ecode, nil
+	}
+
 	var tot time.Duration
 	for tot < cos.DfltRateMaxWait {
-		ecode, err = bp.Backend.DeleteObj(ctx, lom)
-		if err == nil {
-			return ecode, nil
-		}
-		sleep := bp.retry(ctx, err, lom.Bck(), http.MethodDelete, stats.DeleteRateRetryLatencyTotal)
+		sleep := bp.retry(err, lom.Bck(), http.MethodDelete)
 		if sleep == 0 {
-			return ecode, err
+			break
 		}
 		tot += sleep
+		ecode, err = bp.Backend.DeleteObj(ctx, lom)
+		if err == nil {
+			break
+		}
 	}
+	bp.stats(ctx, lom.Bck(), stats.DeleteRateRetryLatencyTotal, tot)
 	return ecode, err
 }
 
 //
 // apply (compare with ais/prate_limit)
 //
 
-func (bp *rlbackend) retry(ctx context.Context, err error, bck *meta.Bck, verb, metric string) time.Duration {
+func (bp *rlbackend) retry(err error, bck *meta.Bck, verb string) time.Duration {
 	if !cmn.IsErrTooManyRequests(err) || !bck.Props.RateLimit.Backend.Enabled {
-		return 0
+		return 0 // not retrying
 	}
 
 	var (
@@ -99,12 +116,18 @@ func (bp *rlbackend) retry(ctx context.Context, err error, bck *meta.Bck, verb,
 	}
 
 	// onerr
-	sleep := arl.OnErr()
+	return arl.OnErr()
+}
 
-	_ = ctx // TODO -- FIXME: in progress
-	vlabs := map[string]string{stats.VarlabBucket: bck.Cname(""), stats.VarlabXactKind: ""}
+func (bp *rlbackend) stats(ctx context.Context, bck *meta.Bck, metric string, total time.Duration) {
+	if total == 0 {
+		return
+	}
+	vlabs := xact.GetCtxVlabs(ctx) // fast path
+	if vlabs == nil {
+		vlabs = map[string]string{stats.VarlabBucket: bck.Cname(""), stats.VarlabXactKind: ""}
+	}
 	bp.t.statsT.AddWith(
-		cos.NamedVal64{Name: metric, Value: int64(sleep), VarLabs: vlabs},
+		cos.NamedVal64{Name: metric, Value: int64(total), VarLabs: vlabs},
 	)
-	return sleep
 }
diff --git a/xact/context.go b/xact/context.go
@@ -0,0 +1,24 @@
+// Package xact provides core functionality for the AIStore eXtended Actions (xactions).
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ */
+package xact
+
+import (
+	"context"
+)
+
+var contextVlabs = &struct{}{}
+
+func NewCtxVlabs(vlabs map[string]string) context.Context {
+	return context.WithValue(context.Background(), contextVlabs, vlabs)
+}
+
+func GetCtxVlabs(ctx context.Context) map[string]string {
+	v := ctx.Value(contextVlabs)
+	vlabs, ok := v.(map[string]string)
+	if !ok {
+		return nil
+	}
+	return vlabs
+}
diff --git a/xact/xs/prefetch.go b/xact/xs/prefetch.go
@@ -50,8 +50,8 @@ type (
 		mu      sync.Mutex
 	}
 	prefetch struct {
-		bp     core.Backend
 		ctx    context.Context
+		bp     core.Backend
 		config *cmn.Config
 		msg    *apc.PrefetchMsg
 		vlabs  map[string]string
@@ -118,7 +118,7 @@ func newPrefetch(xargs *xreg.Args, kind string, bck *meta.Bck, msg *apc.Prefetch
 		stats.VarlabBucket:   bck.Cname(""),
 		stats.VarlabXactKind: r.Kind(),
 	}
-	r.ctx = context.Background()
+	r.ctx = xact.NewCtxVlabs(r.vlabs)
 
 	if r.msg.BlobThreshold > 0 {
 		r.pebl.init(r)
@@ -181,7 +181,6 @@ func (r *prefetch) do(lom *core.LOM, lrit *lrit) {
 	if r.rate.arl != nil {
 		r.rate.arl.RetryAcquire(r.rate.sleep)
 	}
-retry:
 	if r.msg.BlobThreshold > 0 && size >= r.msg.BlobThreshold && !r.pebl.busy() {
 		err = r.blobdl(lom, oa)
 	} else {
@@ -202,12 +201,6 @@ retry:
 		}
 	case cos.IsErrOOS(err):
 		r.Abort(err)
-	case cmn.IsErrTooManyRequests(err):
-		if r.rate.arl != nil {
-			onmanyreq(r.rate.arl, r.vlabs)
-			goto retry
-		}
-		fallthrough
 	default:
 		r.AddErr(err, 5, cos.SmoduleXs)
 	}
@@ -234,7 +227,8 @@ func (r *prefetch) _whinge(lom *core.LOM, size int64) {
 }
 
 // OwtGetPrefetchLock: minimal locking, optimistic concurrency
-// (light-weight alternative to t.GetCold impl.)
+// - light-weight alternative to t.GetCold impl.
+// - rate limited via ais/rlbackend, if defined
 func (r *prefetch) getCold(lom *core.LOM) (ecode int, err error) {
 	started := mono.NanoTime()
 	if ecode, err = r.bp.GetObj(r.ctx, lom, cmn.OwtGetPrefetchLock, nil /*origReq*/); err != nil {
diff --git a/xact/xs/xutils.go b/xact/xs/xutils.go
@@ -9,7 +9,6 @@ import (
 	"time"
 
 	"github.com/NVIDIA/aistore/cmn/cos"
-	"github.com/NVIDIA/aistore/cmn/debug"
 	"github.com/NVIDIA/aistore/cmn/mono"
 	"github.com/NVIDIA/aistore/cmn/nlog"
 	"github.com/NVIDIA/aistore/core"
@@ -27,14 +26,6 @@ func rgetstats(bp core.Backend, vlabs map[string]string, size, started int64) {
 	)
 }
 
-func onmanyreq(arl *cos.AdaptRateLim, vlabs map[string]string) {
-	tstats := core.T.StatsUpdater()
-	sleep := arl.OnErr()
-	tstats.AddWith(
-		cos.NamedVal64{Name: stats.GetRateRetryLatencyTotal, Value: int64(sleep), VarLabs: vlabs},
-	)
-}
-
 ////////////
 // tcrate //
 ////////////
@@ -86,7 +77,12 @@ func newrate(src, dst *meta.Bck, nat int) *tcrate {
 	return &rate
 }
 
+// TODO -- FIXME: remove
 func (rate *tcrate) onmanyreq(vlabs map[string]string) {
-	debug.Assert(rate.onerr != nil)
-	onmanyreq(rate.onerr, vlabs)
+	arl := rate.onerr
+	tstats := core.T.StatsUpdater()
+	sleep := arl.OnErr()
+	tstats.AddWith(
+		cos.NamedVal64{Name: stats.GetRateRetryLatencyTotal, Value: int64(sleep), VarLabs: vlabs},
+	)
 }
