From d20456ba2ff4c5972d568cda46963b7bc4c5d124 Mon Sep 17 00:00:00 2001
From: KT <416432526@qq.com>
Date: Thu, 14 Nov 2024 21:57:50 +0800
Subject: [PATCH] fix runtime error: comparing uncomparable type

Signed-off-by: KT <416432526@qq.com>
---
 clientv3/integration/lease_test.go | 29 ++++++++++++++++++++++++++++-
 clientv3/lease.go                  | 12 +++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/clientv3/integration/lease_test.go b/clientv3/integration/lease_test.go
index 84f18d52905..918b507bace 100644
--- a/clientv3/integration/lease_test.go
+++ b/clientv3/integration/lease_test.go
@@ -133,7 +133,14 @@ func TestLeaseKeepAlive(t *testing.T) {
 		t.Errorf("failed to create lease %v", err)
 	}
 
-	rc, kerr := lapi.KeepAlive(context.Background(), resp.ID)
+	type uncomparableCtx struct {
+		context.Context
+		_ func()
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	rc, kerr := lapi.KeepAlive(uncomparableCtx{Context: ctx}, resp.ID)
 	if kerr != nil {
 		t.Errorf("failed to keepalive lease %v", kerr)
 	}
@@ -151,6 +158,26 @@ func TestLeaseKeepAlive(t *testing.T) {
 		t.Errorf("ID = %x, want %x", kresp.ID, resp.ID)
 	}
 
+	ctx2, cancel2 := context.WithCancel(context.Background())
+	rc2, kerr2 := lapi.KeepAlive(uncomparableCtx{Context: ctx2}, resp.ID)
+	if kerr2 != nil {
+		t.Errorf("failed to keepalive lease %v", kerr2)
+	}
+
+	cancel2()
+
+	_, ok = <-rc2
+	if ok {
+		t.Errorf("chan is not closed, want cancel stop keepalive")
+	}
+
+	select {
+	case <-rc:
+		// cancel2() should not affect first keepalive
+		t.Errorf("chan is closed, want keepalive continue")
+	default:
+	}
+
 	lapi.Close()
 
 	_, ok = <-rc
diff --git a/clientv3/lease.go b/clientv3/lease.go
index cf17f7402f2..cb88e8d1f92 100644
--- a/clientv3/lease.go
+++ b/clientv3/lease.go
@@ -263,6 +263,12 @@ func (l *lessor) Leases(ctx context.Context) (*LeaseLeasesResponse, error) {
 	return nil, toErr(ctx, err)
 }
 
+// To identify the context passed to `KeepAlive`, a key/value pair is
+// attached to the context. The key is a `keepAliveCtxKey` object, and
+// the value is the pointer to the context object itself, ensuring
+// uniqueness as each context has a unique memory address.
+type keepAliveCtxKey struct{}
+
 func (l *lessor) KeepAlive(ctx context.Context, id LeaseID) (<-chan *LeaseKeepAliveResponse, error) {
 	ch := make(chan *LeaseKeepAliveResponse, LeaseResponseChSize)
 
@@ -277,6 +283,10 @@ func (l *lessor) KeepAlive(ctx context.Context, id LeaseID) (<-chan *LeaseKeepAl
 	default:
 	}
 	ka, ok := l.keepAlives[id]
+
+	if ctx.Done() != nil {
+		ctx = context.WithValue(ctx, keepAliveCtxKey{}, &ctx)
+	}
 	if !ok {
 		// create fresh keep alive
 		ka = &keepAlive{
@@ -347,7 +357,7 @@ func (l *lessor) keepAliveCtxCloser(ctx context.Context, id LeaseID, donec <-cha
 
 	// close channel and remove context if still associated with keep alive
 	for i, c := range ka.ctxs {
-		if c == ctx {
+		if c.Value(keepAliveCtxKey{}) == ctx.Value(keepAliveCtxKey{}) {
 			close(ka.chs[i])
 			ka.ctxs = append(ka.ctxs[:i], ka.ctxs[i+1:]...)
 			ka.chs = append(ka.chs[:i], ka.chs[i+1:]...)