From 7263089022e5a1615620b463eb70f4757697a873 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Wed, 25 Apr 2018 10:01:33 -0700
Subject: [PATCH] stickiness: avoid using unsafe (#2023)

App Engine doesn't allow use of unsafe.
---
 picker_wrapper.go | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/picker_wrapper.go b/picker_wrapper.go
index 99b8025e3358..0a984e6c8af0 100644
--- a/picker_wrapper.go
+++ b/picker_wrapper.go
@@ -22,7 +22,6 @@ import (
 	"io"
 	"sync"
 	"sync/atomic"
-	"unsafe"
 
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/balancer"
@@ -47,7 +46,7 @@ type pickerWrapper struct {
 	connErrMu sync.Mutex
 	connErr   error
 
-	stickinessMDKey unsafe.Pointer
+	stickinessMDKey atomic.Value
 	stickiness      *stickyStore
 }
 
@@ -73,18 +72,17 @@ func (bp *pickerWrapper) connectionError() error {
 }
 
 func (bp *pickerWrapper) updateStickinessMDKey(newKey string) {
-	oldKeyUnsafePtr := atomic.SwapPointer(&bp.stickinessMDKey, unsafe.Pointer(&newKey))
-	if ptr := (*string)(oldKeyUnsafePtr); ptr == nil || *ptr != newKey {
+	// No need to check ok because mdKey == "" if ok == false.
+	if oldKey, _ := bp.stickinessMDKey.Load().(string); oldKey != newKey {
+		bp.stickinessMDKey.Store(newKey)
 		bp.stickiness.reset(newKey)
 	}
 }
 
 func (bp *pickerWrapper) getStickinessMDKey() string {
-	mdKeyPtr := (*string)(atomic.LoadPointer(&bp.stickinessMDKey))
-	if mdKeyPtr != nil {
-		return *mdKeyPtr
-	}
-	return ""
+	// No need to check ok because mdKey == "" if ok == false.
+	mdKey, _ := bp.stickinessMDKey.Load().(string)
+	return mdKey
 }
 
 func (bp *pickerWrapper) clearStickinessState() {