@@ -99,7 +99,6 @@ type asyncClient struct {
9999 containerID uint64
100100 adapter metadata.Adapter
101101 id uint64
102- contexts sync.Map // id -> request
103102 leaderConn goetty.IOSession
104103
105104 resetReadC chan string
@@ -113,6 +112,11 @@ type asyncClient struct {
113112 sync.RWMutex
114113 state int32
115114 }
115+
116+ contextsMu struct {
117+ sync.RWMutex
118+ contexts map [uint64 ]* ctx
119+ }
116120}
117121
118122// NewClient create a prophet client
@@ -132,6 +136,7 @@ func NewClient(adapter metadata.Adapter, opts ...Option) Client {
132136 c .opts .adjust ()
133137 c .stopper = stop .NewStopper ("prophet-client" , stop .WithLogger (c .opts .logger ))
134138 c .leaderConn = createConn (c .opts .logger )
139+ c .contextsMu .contexts = make (map [uint64 ]* ctx )
135140 c .start ()
136141 return c
137142}
@@ -592,8 +597,10 @@ func (c *asyncClient) do(ctx *ctx) error {
592597 if ! added {
593598 ctx .req .ID = c .nextID ()
594599 if ctx .sync || ctx .cb != nil {
595- c .contexts .Store (ctx .req .ID , ctx )
600+ c .contextsMu .Lock ()
601+ c .contextsMu .contexts [ctx .req .ID ] = ctx
596602 util .DefaultTimeoutWheel ().Schedule (c .opts .rpcTimeout , c .timeout , ctx .req .ID )
603+ c .contextsMu .Unlock ()
597604 }
598605 added = true
599606 }
@@ -612,9 +619,11 @@ func (c *asyncClient) do(ctx *ctx) error {
612619}
613620
614621func (c * asyncClient ) timeout (arg interface {}) {
615- if v , ok := c .contexts .Load (arg ); ok {
616- c .contexts .Delete (arg )
617- v .(* ctx ).done (nil , ErrTimeout )
622+ c .contextsMu .RLock ()
623+ defer c .contextsMu .RUnlock ()
624+
625+ if ctx , ok := c .contextsMu .contexts [arg .(uint64 )]; ok {
626+ ctx .done (nil , ErrTimeout )
618627 }
619628}
620629
@@ -667,12 +676,12 @@ OUTER:
667676 for {
668677 select {
669678 case <- stopCtx .Done ():
670- c .contexts . Range ( func ( key , value interface {}) bool {
671- if value != nil {
672- value .( * ctx ) .done (nil , ErrClosed )
673- }
674- return true
675- } )
679+ c .contextsMu . Lock ()
680+ for k , ctx := range c . contextsMu . contexts {
681+ ctx .done (nil , ErrClosed )
682+ delete ( c . contextsMu . contexts , k )
683+ }
684+ c . contextsMu . Unlock ( )
676685 return
677686 case leader , ok := <- c .resetReadC :
678687 if ok {
@@ -714,20 +723,25 @@ OUTER:
714723}
715724
716725func (c * asyncClient ) requestDoneWithRetry (resp * rpcpb.Response ) {
717- v , ok := c .contexts .Load (resp .ID )
718- if ok && v != nil {
719- v .(* ctx ).done (nil , util .ErrNotLeader )
726+ c .contextsMu .Lock ()
727+ defer c .contextsMu .Unlock ()
728+
729+ if ctx , ok := c .contextsMu .contexts [resp .ID ]; ok {
730+ delete (c .contextsMu .contexts , resp .ID )
731+ ctx .done (nil , util .ErrNotLeader )
720732 }
721733}
722734
723735func (c * asyncClient ) requestDone (resp * rpcpb.Response ) {
724- v , ok := c .contexts .Load (resp .ID )
725- if ok && v != nil {
726- c .contexts .Delete (resp .ID )
736+ c .contextsMu .Lock ()
737+ defer c .contextsMu .Unlock ()
738+
739+ if ctx , ok := c .contextsMu .contexts [resp .ID ]; ok {
740+ delete (c .contextsMu .contexts , resp .ID )
727741 if resp .Error != "" {
728- v .( * ctx ) .done (nil , errors .New (resp .Error ))
742+ ctx .done (nil , errors .New (resp .Error ))
729743 } else {
730- v .( * ctx ) .done (resp , nil )
744+ ctx .done (resp , nil )
731745 }
732746 }
733747}
0 commit comments