github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/lbcache/hookable.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package lbcache
    18  
    19  import (
    20  	"sync"
    21  
    22  	"github.com/cloudwego/kitex/pkg/discovery"
    23  	"github.com/cloudwego/kitex/pkg/loadbalance"
    24  )
    25  
    26  // intercept the rebalancer and call hooks
    27  // wrap the loadbalance.Rebalancer and execute registered hooks
    28  type hookableRebalancer struct {
    29  	inner          loadbalance.Rebalancer
    30  	rebalanceL     sync.Mutex
    31  	rebalanceIndex int
    32  	rebalanceHooks map[int]func(*discovery.Change)
    33  	deleteL        sync.Mutex
    34  	deleteIndex    int
    35  	deleteHooks    map[int]func(*discovery.Change)
    36  }
    37  
    38  var (
    39  	_ loadbalance.Rebalancer = (*hookableRebalancer)(nil)
    40  	_ Hookable               = (*hookableRebalancer)(nil)
    41  )
    42  
    43  func newHookRebalancer(inner loadbalance.Rebalancer) *hookableRebalancer {
    44  	return &hookableRebalancer{
    45  		inner:          inner,
    46  		rebalanceHooks: map[int]func(*discovery.Change){},
    47  		deleteHooks:    map[int]func(*discovery.Change){},
    48  	}
    49  }
    50  
    51  func (b *hookableRebalancer) Rebalance(ch discovery.Change) {
    52  	b.rebalanceL.Lock()
    53  	for _, h := range b.rebalanceHooks {
    54  		h(&ch)
    55  	}
    56  	b.rebalanceL.Unlock()
    57  
    58  	b.inner.Rebalance(ch)
    59  }
    60  
    61  func (b *hookableRebalancer) Delete(ch discovery.Change) {
    62  	b.deleteL.Lock()
    63  	for _, h := range b.deleteHooks {
    64  		h(&ch)
    65  	}
    66  	b.deleteL.Unlock()
    67  
    68  	b.inner.Delete(ch)
    69  }
    70  
    71  func (b *hookableRebalancer) RegisterRebalanceHook(f func(ch *discovery.Change)) int {
    72  	b.rebalanceL.Lock()
    73  	defer b.rebalanceL.Unlock()
    74  
    75  	index := b.rebalanceIndex
    76  	b.rebalanceIndex++
    77  	b.rebalanceHooks[index] = f
    78  	return index
    79  }
    80  
    81  func (b *hookableRebalancer) DeregisterRebalanceHook(index int) {
    82  	b.rebalanceL.Lock()
    83  	defer b.rebalanceL.Unlock()
    84  
    85  	delete(b.rebalanceHooks, index)
    86  }
    87  
    88  func (b *hookableRebalancer) RegisterDeleteHook(f func(ch *discovery.Change)) int {
    89  	b.deleteL.Lock()
    90  	defer b.deleteL.Unlock()
    91  
    92  	index := b.deleteIndex
    93  	b.deleteIndex++
    94  	b.deleteHooks[index] = f
    95  	return index
    96  }
    97  
    98  func (b *hookableRebalancer) DeregisterDeleteHook(index int) {
    99  	b.deleteL.Lock()
   100  	defer b.deleteL.Unlock()
   101  
   102  	delete(b.deleteHooks, index)
   103  }
   104  
   105  type noopHookRebalancer struct{}
   106  
   107  var _ Hookable = (*noopHookRebalancer)(nil)
   108  
   109  func (noopHookRebalancer) RegisterRebalanceHook(func(ch *discovery.Change)) int { return 0 }
   110  func (noopHookRebalancer) RegisterDeleteHook(func(ch *discovery.Change)) int    { return 0 }
   111  func (noopHookRebalancer) DeregisterRebalanceHook(index int)                    {}
   112  func (noopHookRebalancer) DeregisterDeleteHook(index int)                       {}