trpc.group/trpc-go/trpc-go@v1.0.3/naming/bannednodes/banned_nodes.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package bannednodes defines a concurrent safe node list which may be bound to a context.
    15  package bannednodes
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  
    21  	"trpc.group/trpc-go/trpc-go/naming/registry"
    22  )
    23  
    24  // ctxKeyBannedNodes is the key of the context.
    25  type ctxKeyBannedNodes struct{}
    26  
    27  // bannedNodes is the value of the context.
    28  type bannedNodes struct {
    29  	mu        sync.Mutex
    30  	nodes     *Nodes
    31  	mandatory bool
    32  }
    33  
    34  // NewCtx creates a new context and sets its k-v.
    35  func NewCtx(ctx context.Context, mandatory bool) context.Context {
    36  	return context.WithValue(ctx, ctxKeyBannedNodes{}, &bannedNodes{mandatory: mandatory})
    37  }
    38  
    39  // FromCtx returns the node list and a boolean which indicate whether it is abandoned.
    40  // FromCtx does not return a bannedNodes, but a read only linked list. The internal lock is not
    41  // exported to external user.
    42  func FromCtx(ctx context.Context) (nodes *Nodes, mandatory bool, ok bool) {
    43  	bannedNodes, ok := ctx.Value(ctxKeyBannedNodes{}).(*bannedNodes)
    44  	if !ok {
    45  		return nil, false, false
    46  	}
    47  
    48  	bannedNodes.mu.Lock()
    49  	defer bannedNodes.mu.Unlock()
    50  	return bannedNodes.nodes, bannedNodes.mandatory, true
    51  }
    52  
    53  // Add adds a new node to ctx.
    54  // Nothing would happen if there's no k-v.
    55  func Add(ctx context.Context, nodes ...*registry.Node) {
    56  	bannedNodes, ok := ctx.Value(ctxKeyBannedNodes{}).(*bannedNodes)
    57  	if !ok {
    58  		return
    59  	}
    60  
    61  	bannedNodes.mu.Lock()
    62  	defer bannedNodes.mu.Unlock()
    63  	for _, node := range nodes {
    64  		bannedNodes.nodes = &Nodes{
    65  			next: bannedNodes.nodes,
    66  			node: node,
    67  		}
    68  	}
    69  }
    70  
    71  // Nodes is a linked list of registry.Node.
    72  type Nodes struct {
    73  	next *Nodes
    74  	node *registry.Node
    75  }
    76  
    77  // Range likes the Range method of sync.Map.
    78  // It calls f serially for each node in Nodes. Range stops on f failed.
    79  // Range returns true after traversing all nodes, false otherwise.
    80  // Users should not change n in f.
    81  func (nodes *Nodes) Range(f func(n *registry.Node) bool) bool {
    82  	if nodes == nil {
    83  		return true
    84  	}
    85  	if f(nodes.node) {
    86  		return nodes.next.Range(f)
    87  	}
    88  	return false
    89  }