go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/cv/internal/buildbucket/fake/timed_map.go (about)

     1  // Copyright 2022 The LUCI Authors.
     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 bbfake
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"time"
    21  
    22  	"go.chromium.org/luci/common/clock"
    23  )
    24  
    25  type item struct {
    26  	value    any
    27  	expireAt time.Time
    28  }
    29  
    30  // timedMap is a map where each item has a TTL.
    31  //
    32  // cleanup is done lazily.
    33  type timedMap struct {
    34  	m  map[string]*item
    35  	mu sync.RWMutex
    36  }
    37  
    38  // set adds an item to the map with ttl.
    39  //
    40  // zero or negative `expiration` means never expire
    41  func (rc *timedMap) set(ctx context.Context, key string, value any, expiration time.Duration) {
    42  	rc.mu.Lock()
    43  	defer rc.mu.Unlock()
    44  	i := &item{value: value}
    45  	if expiration > 0 {
    46  		i.expireAt = clock.Now(ctx).Add(expiration)
    47  	}
    48  	if rc.m == nil {
    49  		rc.m = make(map[string]*item)
    50  	}
    51  	rc.m[key] = i
    52  	rc.cleanupLocked(ctx)
    53  }
    54  
    55  func (rc *timedMap) get(ctx context.Context, key string) (any, bool) {
    56  	defer func() {
    57  		// TODO(yiwzhang): Once move to go 1.18, use Trylock for best effort clean
    58  		// up.
    59  		rc.mu.Lock()
    60  		rc.cleanupLocked(ctx)
    61  		rc.mu.Unlock()
    62  	}()
    63  	rc.mu.RLock()
    64  	defer rc.mu.RUnlock()
    65  	switch item, ok := rc.m[key]; {
    66  	case !ok:
    67  		return nil, false
    68  	case item.expired(clock.Now(ctx)):
    69  		return nil, false
    70  	default:
    71  		return item.value, true
    72  	}
    73  }
    74  
    75  func (rc *timedMap) cleanupLocked(ctx context.Context) {
    76  	now := clock.Now(ctx)
    77  	for key, item := range rc.m {
    78  		if item.expired(now) {
    79  			delete(rc.m, key)
    80  		}
    81  	}
    82  }
    83  
    84  func (item *item) expired(now time.Time) bool {
    85  	return !item.expireAt.IsZero() && item.expireAt.Before(now)
    86  }