vitess.io/vitess@v0.16.2/go/sync2/consolidator.go (about)

     1  /*
     2  Copyright 2019 The Vitess 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 sync2
    18  
    19  import (
    20  	"sync"
    21  	"sync/atomic"
    22  
    23  	"vitess.io/vitess/go/cache"
    24  )
    25  
    26  // Consolidator consolidates duplicate queries from executing simulaneously
    27  // and shares results between them.
    28  type Consolidator struct {
    29  	*ConsolidatorCache
    30  
    31  	mu      sync.Mutex
    32  	queries map[string]*Result
    33  }
    34  
    35  // NewConsolidator creates a new Consolidator
    36  func NewConsolidator() *Consolidator {
    37  	return &Consolidator{
    38  		queries:           make(map[string]*Result),
    39  		ConsolidatorCache: NewConsolidatorCache(1000),
    40  	}
    41  }
    42  
    43  // Result is a wrapper for result of a query.
    44  type Result struct {
    45  	// executing is used to block additional requests.
    46  	// The original request holds a write lock while additional ones are blocked
    47  	// on acquiring a read lock (see Wait() below.)
    48  	executing    sync.RWMutex
    49  	consolidator *Consolidator
    50  	query        string
    51  	Result       any
    52  	Err          error
    53  }
    54  
    55  // Create adds a query to currently executing queries and acquires a
    56  // lock on its Result if it is not already present. If the query is
    57  // a duplicate, Create returns false.
    58  func (co *Consolidator) Create(query string) (r *Result, created bool) {
    59  	co.mu.Lock()
    60  	defer co.mu.Unlock()
    61  	if r, ok := co.queries[query]; ok {
    62  		return r, false
    63  	}
    64  	r = &Result{consolidator: co, query: query}
    65  	r.executing.Lock()
    66  	co.queries[query] = r
    67  	return r, true
    68  }
    69  
    70  // Broadcast removes the entry from current queries and releases the
    71  // lock on its Result. Broadcast should be invoked when original
    72  // query completes execution.
    73  func (rs *Result) Broadcast() {
    74  	rs.consolidator.mu.Lock()
    75  	defer rs.consolidator.mu.Unlock()
    76  	delete(rs.consolidator.queries, rs.query)
    77  	rs.executing.Unlock()
    78  }
    79  
    80  // Wait waits for the original query to complete execution. Wait should
    81  // be invoked for duplicate queries.
    82  func (rs *Result) Wait() {
    83  	rs.consolidator.Record(rs.query)
    84  	rs.executing.RLock()
    85  }
    86  
    87  // ConsolidatorCache is a thread-safe object used for counting how often recent
    88  // queries have been consolidated.
    89  // It is also used by the txserializer package to count how often transactions
    90  // have been queued and had to wait because they targeted the same row (range).
    91  type ConsolidatorCache struct {
    92  	*cache.LRUCache
    93  }
    94  
    95  // NewConsolidatorCache creates a new cache with the given capacity.
    96  func NewConsolidatorCache(capacity int64) *ConsolidatorCache {
    97  	return &ConsolidatorCache{cache.NewLRUCache(capacity, func(_ any) int64 {
    98  		return 1
    99  	})}
   100  }
   101  
   102  // Record increments the count for "query" by 1.
   103  // If it's not in the cache yet, it will be added.
   104  func (cc *ConsolidatorCache) Record(query string) {
   105  	if v, ok := cc.Get(query); ok {
   106  		v.(*ccount).add(1)
   107  	} else {
   108  		c := ccount(1)
   109  		cc.Set(query, &c)
   110  	}
   111  }
   112  
   113  // ConsolidatorCacheItem is a wrapper for the items in the consolidator cache
   114  type ConsolidatorCacheItem struct {
   115  	Query string
   116  	Count int64
   117  }
   118  
   119  // Items returns the items in the cache as an array of String, int64 structs
   120  func (cc *ConsolidatorCache) Items() []ConsolidatorCacheItem {
   121  	items := cc.LRUCache.Items()
   122  	ret := make([]ConsolidatorCacheItem, len(items))
   123  	for i, v := range items {
   124  		ret[i] = ConsolidatorCacheItem{Query: v.Key, Count: v.Value.(*ccount).get()}
   125  	}
   126  	return ret
   127  }
   128  
   129  // ccount elements are used with a cache.LRUCache object to track if another
   130  // request for the same query is already in progress.
   131  type ccount int64
   132  
   133  func (cc *ccount) add(n int64) int64 {
   134  	return atomic.AddInt64((*int64)(cc), n)
   135  }
   136  
   137  func (cc *ccount) get() int64 {
   138  	return atomic.LoadInt64((*int64)(cc))
   139  }