github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/catalog/lease/helpers_test.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package lease
    12  
    13  import (
    14  	"time"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    17  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    18  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    19  )
    20  
    21  // A unique id for a particular table descriptor version.
    22  type tableVersionID struct {
    23  	id      sqlbase.ID
    24  	version sqlbase.DescriptorVersion
    25  }
    26  
    27  // LeaseRemovalTracker can be used to wait for leases to be removed from the
    28  // store (leases are removed from the store async w.r.t. Manager
    29  // operations).
    30  // To use it, its LeaseRemovedNotification method must be hooked up to
    31  // StorageTestingKnobs.LeaseReleasedEvent. Then, every time you want to wait
    32  // for a lease, get a tracker object through TrackRemoval() before calling
    33  // Manager.Release(), and then call WaitForRemoval() on the tracker to
    34  // block for the removal from the store.
    35  //
    36  // All methods are thread-safe.
    37  type LeaseRemovalTracker struct {
    38  	mu syncutil.Mutex
    39  	// map from a lease whose release we're waiting for to a tracker for that
    40  	// lease.
    41  	tracking map[tableVersionID]RemovalTracker
    42  }
    43  
    44  type RemovalTracker struct {
    45  	removed chan struct{}
    46  	// Pointer to a shared err. *err is written when removed is closed.
    47  	err *error
    48  }
    49  
    50  // NewLeaseRemovalTracker creates a LeaseRemovalTracker.
    51  func NewLeaseRemovalTracker() *LeaseRemovalTracker {
    52  	return &LeaseRemovalTracker{
    53  		tracking: make(map[tableVersionID]RemovalTracker),
    54  	}
    55  }
    56  
    57  // TrackRemoval starts monitoring lease removals for a particular lease.
    58  // This should be called before triggering the operation that (asynchronously)
    59  // removes the lease.
    60  func (w *LeaseRemovalTracker) TrackRemoval(table *sqlbase.ImmutableTableDescriptor) RemovalTracker {
    61  	id := tableVersionID{
    62  		id:      table.ID,
    63  		version: table.Version,
    64  	}
    65  	w.mu.Lock()
    66  	defer w.mu.Unlock()
    67  	if tracker, ok := w.tracking[id]; ok {
    68  		return tracker
    69  	}
    70  	tracker := RemovalTracker{removed: make(chan struct{}), err: new(error)}
    71  	w.tracking[id] = tracker
    72  	return tracker
    73  }
    74  
    75  // WaitForRemoval blocks until the lease is removed from the store.
    76  func (t RemovalTracker) WaitForRemoval() error {
    77  	<-t.removed
    78  	return *t.err
    79  }
    80  
    81  // LeaseRemovedNotification has to be called after a lease is removed from the
    82  // store. This should be hooked up as a callback to
    83  // StorageTestingKnobs.LeaseReleasedEvent.
    84  func (w *LeaseRemovalTracker) LeaseRemovedNotification(
    85  	id sqlbase.ID, version sqlbase.DescriptorVersion, err error,
    86  ) {
    87  	w.mu.Lock()
    88  	defer w.mu.Unlock()
    89  
    90  	idx := tableVersionID{
    91  		id:      id,
    92  		version: version,
    93  	}
    94  
    95  	if tracker, ok := w.tracking[idx]; ok {
    96  		*tracker.err = err
    97  		close(tracker.removed)
    98  		delete(w.tracking, idx)
    99  	}
   100  }
   101  
   102  func (m *Manager) ExpireLeases(clock *hlc.Clock) {
   103  	past := clock.Now().GoTime().Add(-time.Millisecond)
   104  
   105  	m.tableNames.mu.Lock()
   106  	for _, table := range m.tableNames.tables {
   107  		table.expiration = hlc.Timestamp{WallTime: past.UnixNano()}
   108  	}
   109  	m.tableNames.mu.Unlock()
   110  }