gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/internal/cache/timeoutCache.go (about)

     1  /*
     2   * Copyright 2019 gRPC 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 cache implements caches to be used in gRPC.
    18  package cache
    19  
    20  import (
    21  	"sync"
    22  	"time"
    23  )
    24  
    25  type cacheEntry struct {
    26  	item interface{}
    27  	// Note that to avoid deadlocks (potentially caused by lock ordering),
    28  	// callback can only be called without holding cache's mutex.
    29  	callback func()
    30  	timer    *time.Timer
    31  	// deleted is set to true in Remove() when the call to timer.Stop() fails.
    32  	// This can happen when the timer in the cache entry fires around the same
    33  	// time that timer.stop() is called in Remove().
    34  	deleted bool
    35  }
    36  
    37  // TimeoutCache is a cache with items to be deleted after a timeout.
    38  type TimeoutCache struct {
    39  	mu      sync.Mutex
    40  	timeout time.Duration
    41  	cache   map[interface{}]*cacheEntry
    42  }
    43  
    44  // NewTimeoutCache creates a TimeoutCache with the given timeout.
    45  func NewTimeoutCache(timeout time.Duration) *TimeoutCache {
    46  	return &TimeoutCache{
    47  		timeout: timeout,
    48  		cache:   make(map[interface{}]*cacheEntry),
    49  	}
    50  }
    51  
    52  // Add adds an item to the cache, with the specified callback to be called when
    53  // the item is removed from the cache upon timeout. If the item is removed from
    54  // the cache using a call to Remove before the timeout expires, the callback
    55  // will not be called.
    56  //
    57  // If the Add was successful, it returns (newly added item, true). If there is
    58  // an existing entry for the specified key, the cache entry is not be updated
    59  // with the specified item and it returns (existing item, false).
    60  func (c *TimeoutCache) Add(key, item interface{}, callback func()) (interface{}, bool) {
    61  	c.mu.Lock()
    62  	defer c.mu.Unlock()
    63  	if e, ok := c.cache[key]; ok {
    64  		return e.item, false
    65  	}
    66  
    67  	entry := &cacheEntry{
    68  		item:     item,
    69  		callback: callback,
    70  	}
    71  	entry.timer = time.AfterFunc(c.timeout, func() {
    72  		c.mu.Lock()
    73  		if entry.deleted {
    74  			c.mu.Unlock()
    75  			// Abort the delete since this has been taken care of in Remove().
    76  			return
    77  		}
    78  		delete(c.cache, key)
    79  		c.mu.Unlock()
    80  		entry.callback()
    81  	})
    82  	c.cache[key] = entry
    83  	return item, true
    84  }
    85  
    86  // Remove the item with the key from the cache.
    87  //
    88  // If the specified key exists in the cache, it returns (item associated with
    89  // key, true) and the callback associated with the item is guaranteed to be not
    90  // called. If the given key is not found in the cache, it returns (nil, false)
    91  func (c *TimeoutCache) Remove(key interface{}) (item interface{}, ok bool) {
    92  	c.mu.Lock()
    93  	defer c.mu.Unlock()
    94  	entry, ok := c.removeInternal(key)
    95  	if !ok {
    96  		return nil, false
    97  	}
    98  	return entry.item, true
    99  }
   100  
   101  // removeInternal removes and returns the item with key.
   102  //
   103  // caller must hold c.mu.
   104  func (c *TimeoutCache) removeInternal(key interface{}) (*cacheEntry, bool) {
   105  	entry, ok := c.cache[key]
   106  	if !ok {
   107  		return nil, false
   108  	}
   109  	delete(c.cache, key)
   110  	if !entry.timer.Stop() {
   111  		// If stop was not successful, the timer has fired (this can only happen
   112  		// in a race). But the deleting function is blocked on c.mu because the
   113  		// mutex was held by the caller of this function.
   114  		//
   115  		// Set deleted to true to abort the deleting function. When the lock is
   116  		// released, the delete function will acquire the lock, check the value
   117  		// of deleted and return.
   118  		entry.deleted = true
   119  	}
   120  	return entry, true
   121  }
   122  
   123  // Clear removes all entries, and runs the callbacks if runCallback is true.
   124  func (c *TimeoutCache) Clear(runCallback bool) {
   125  	var entries []*cacheEntry
   126  	c.mu.Lock()
   127  	for key := range c.cache {
   128  		if e, ok := c.removeInternal(key); ok {
   129  			entries = append(entries, e)
   130  		}
   131  	}
   132  	c.mu.Unlock()
   133  
   134  	if !runCallback {
   135  		return
   136  	}
   137  
   138  	// removeInternal removes entries from cache, and also stops the timer, so
   139  	// the callback is guaranteed to be not called. If runCallback is true,
   140  	// manual execute all callbacks.
   141  	for _, entry := range entries {
   142  		entry.callback()
   143  	}
   144  }