github.com/matrixorigin/matrixone@v1.2.0/pkg/hakeeper/task/cn_pool.go (about)

     1  // Copyright 2022 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package task
    16  
    17  import (
    18  	"container/heap"
    19  	pb "github.com/matrixorigin/matrixone/pkg/pb/logservice"
    20  	"slices"
    21  )
    22  
    23  type cnStore struct {
    24  	uuid string
    25  	info pb.CNStoreInfo
    26  }
    27  
    28  type cnPool struct {
    29  	freq     map[string]uint32
    30  	sortedCN []cnStore
    31  }
    32  
    33  func (p *cnPool) Len() int {
    34  	return len(p.sortedCN)
    35  }
    36  
    37  func (p *cnPool) Less(i, j int) bool {
    38  	return p.freq[p.sortedCN[i].uuid] < p.freq[p.sortedCN[j].uuid]
    39  }
    40  
    41  func (p *cnPool) Swap(i, j int) {
    42  	p.sortedCN[i], p.sortedCN[j] = p.sortedCN[j], p.sortedCN[i]
    43  }
    44  
    45  func (p *cnPool) Push(x any) {
    46  	store := x.(cnStore)
    47  	p.set(store, p.getFreq(store.uuid)+1)
    48  }
    49  
    50  func (p *cnPool) Pop() any {
    51  	x := p.sortedCN[len(p.sortedCN)-1]
    52  	p.remove(x.uuid)
    53  	return x
    54  }
    55  
    56  func newCNPool() *cnPool {
    57  	pool := &cnPool{
    58  		freq:     make(map[string]uint32),
    59  		sortedCN: make([]cnStore, 0),
    60  	}
    61  	heap.Init(pool)
    62  	return pool
    63  }
    64  
    65  func newCNPoolWithCNState(cnState pb.CNState) *cnPool {
    66  	orderedMap := &cnPool{
    67  		freq:     make(map[string]uint32, len(cnState.Stores)),
    68  		sortedCN: make([]cnStore, 0, len(cnState.Stores)),
    69  	}
    70  
    71  	for key, info := range cnState.Stores {
    72  		orderedMap.freq[key] = 0
    73  		orderedMap.sortedCN = append(orderedMap.sortedCN, cnStore{key, info})
    74  	}
    75  	heap.Init(orderedMap)
    76  	return orderedMap
    77  }
    78  
    79  func (p *cnPool) set(key cnStore, val uint32) {
    80  	if _, ok := p.freq[key.uuid]; !ok {
    81  		p.sortedCN = append(p.sortedCN, key)
    82  	}
    83  	p.freq[key.uuid] = val
    84  	heap.Fix(p, slices.IndexFunc(p.sortedCN,
    85  		func(store cnStore) bool {
    86  			return store.uuid == key.uuid
    87  		}))
    88  }
    89  
    90  func (p *cnPool) getFreq(key string) uint32 {
    91  	return p.freq[key]
    92  }
    93  
    94  func (p *cnPool) contains(uuid string) bool {
    95  	_, ok := p.freq[uuid]
    96  	return ok
    97  }
    98  
    99  func (p *cnPool) min() cnStore {
   100  	if len(p.sortedCN) == 0 {
   101  		return cnStore{}
   102  	}
   103  	return p.sortedCN[0]
   104  }
   105  
   106  func (p *cnPool) remove(key string) {
   107  	delete(p.freq, key)
   108  	p.sortedCN = slices.DeleteFunc(p.sortedCN, func(store cnStore) bool {
   109  		return store.uuid == key
   110  	})
   111  }
   112  
   113  func (p *cnPool) getStore(key string) (cnStore, bool) {
   114  	if _, ok := p.freq[key]; !ok {
   115  		return cnStore{}, ok
   116  	}
   117  
   118  	for _, store := range p.sortedCN {
   119  		if store.uuid == key {
   120  			return store, true
   121  		}
   122  	}
   123  	return cnStore{}, false
   124  }