From 16a5d0b6f2cf7b96f66525ce0dab80b6c35305f6 Mon Sep 17 00:00:00 2001
From: HeyJavaBean <liyun.339@bytedance.com>
Date: Fri, 25 Oct 2024 15:49:45 +0800
Subject: [PATCH] feat: support fallback by ctx

---
 client/client.go           |  6 +++++-
 client/context_fallback.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 client/context_fallback.go

diff --git a/client/client.go b/client/client.go
index c68b2d386c..4a603d53a5 100644
--- a/client/client.go
+++ b/client/client.go
@@ -380,7 +380,11 @@ func (kc *kClient) Call(ctx context.Context, method string, request, response in
 	}
 
 	// do fallback if with setup
-	err, reportErr = doFallbackIfNeeded(ctx, ri, request, response, err, kc.opt.Fallback, callOpts)
+	fb := kc.opt.Fallback
+	if fb == nil {
+		fb = getContextFallback(ctx)
+	}
+	err, reportErr = doFallbackIfNeeded(ctx, ri, request, response, err, fb, callOpts)
 	return err
 }
 
diff --git a/client/context_fallback.go b/client/context_fallback.go
new file mode 100644
index 0000000000..635dad19ec
--- /dev/null
+++ b/client/context_fallback.go
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2022 CloudWeGo Authors
+ *
+ * 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 client
+
+import (
+	"context"
+
+	"github.com/cloudwego/kitex/pkg/fallback"
+)
+
+type ctxFallbackKey struct{}
+
+// WithContextFallback add a fallback into current ctx
+// Every client receive this ctx will execute the fallback.
+// Note that ContextFallback is prior to the ClientFallback.
+func WithContextFallback(ctx context.Context, fb *fallback.Policy) context.Context {
+	return context.WithValue(ctx, ctxFallbackKey{}, fb)
+}
+
+func getContextFallback(ctx context.Context) *fallback.Policy {
+	fb, _ := ctx.Value(ctxFallbackKey{}).(*fallback.Policy)
+	return fb
+}