github.com/lingyao2333/mo-zero@v1.4.1/core/discov/internal/statewatcher.go (about)

     1  //go:generate mockgen -package internal -destination statewatcher_mock.go -source statewatcher.go etcdConn
     2  
     3  package internal
     4  
     5  import (
     6  	"context"
     7  	"sync"
     8  
     9  	"google.golang.org/grpc/connectivity"
    10  )
    11  
    12  type (
    13  	etcdConn interface {
    14  		GetState() connectivity.State
    15  		WaitForStateChange(ctx context.Context, sourceState connectivity.State) bool
    16  	}
    17  
    18  	stateWatcher struct {
    19  		disconnected bool
    20  		currentState connectivity.State
    21  		listeners    []func()
    22  		// lock only guards listeners, because only listens can be accessed by other goroutines.
    23  		lock sync.Mutex
    24  	}
    25  )
    26  
    27  func newStateWatcher() *stateWatcher {
    28  	return new(stateWatcher)
    29  }
    30  
    31  func (sw *stateWatcher) addListener(l func()) {
    32  	sw.lock.Lock()
    33  	sw.listeners = append(sw.listeners, l)
    34  	sw.lock.Unlock()
    35  }
    36  
    37  func (sw *stateWatcher) notifyListeners() {
    38  	sw.lock.Lock()
    39  	defer sw.lock.Unlock()
    40  
    41  	for _, l := range sw.listeners {
    42  		l()
    43  	}
    44  }
    45  
    46  func (sw *stateWatcher) updateState(conn etcdConn) {
    47  	sw.currentState = conn.GetState()
    48  	switch sw.currentState {
    49  	case connectivity.TransientFailure, connectivity.Shutdown:
    50  		sw.disconnected = true
    51  	case connectivity.Ready:
    52  		if sw.disconnected {
    53  			sw.disconnected = false
    54  			sw.notifyListeners()
    55  		}
    56  	}
    57  }
    58  
    59  func (sw *stateWatcher) watch(conn etcdConn) {
    60  	sw.currentState = conn.GetState()
    61  	for {
    62  		if conn.WaitForStateChange(context.Background(), sw.currentState) {
    63  			sw.updateState(conn)
    64  		}
    65  	}
    66  }