gate_unsafe.go raw

   1  // Copyright 2018 The gVisor 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 sync
  16  
  17  import (
  18  	"fmt"
  19  	"math"
  20  	"sync/atomic"
  21  	"unsafe"
  22  
  23  	"gvisor.dev/gvisor/pkg/gohacks"
  24  )
  25  
  26  // Gate is a synchronization primitive that allows concurrent goroutines to
  27  // "enter" it as long as it hasn't been closed yet. Once it's been closed,
  28  // goroutines cannot enter it anymore, but are allowed to leave, and the closer
  29  // will be informed when all goroutines have left.
  30  //
  31  // Gate is similar to WaitGroup:
  32  //
  33  //   - Gate.Enter() is analogous to WaitGroup.Add(1), but may be called even if
  34  //     the Gate counter is 0 and fails if Gate.Close() has been called.
  35  //
  36  //   - Gate.Leave() is equivalent to WaitGroup.Done().
  37  //
  38  //   - Gate.Close() is analogous to WaitGroup.Wait(), but also causes future
  39  //
  40  // calls to Gate.Enter() to fail and may only be called once, from a single
  41  // goroutine.
  42  //
  43  // This is useful, for example, in cases when a goroutine is trying to clean up
  44  // an object for which multiple goroutines have pointers. In such a case, users
  45  // would be required to enter and leave the Gate, and the cleaner would wait
  46  // until all users are gone (and no new ones are allowed) before proceeding.
  47  //
  48  // Users:
  49  //
  50  //	if !g.Enter() {
  51  //		// Gate is closed, we can't use the object.
  52  //		return
  53  //	}
  54  //
  55  //	// Do something with object.
  56  //	[...]
  57  //
  58  //	g.Leave()
  59  //
  60  // Closer:
  61  //
  62  //	// Prevent new users from using the object, and wait for the existing
  63  //	// ones to complete.
  64  //	g.Close()
  65  //
  66  //	// Clean up the object.
  67  //	[...]
  68  type Gate struct {
  69  	userCount int32
  70  	closingG  uintptr
  71  }
  72  
  73  const preparingG = 1
  74  
  75  // Enter tries to enter the gate. It will succeed if it hasn't been closed yet,
  76  // in which case the caller must eventually call Leave().
  77  //
  78  // This function is thread-safe.
  79  func (g *Gate) Enter() bool {
  80  	if atomic.AddInt32(&g.userCount, 1) > 0 {
  81  		return true
  82  	}
  83  	g.leaveAfterFailedEnter()
  84  	return false
  85  }
  86  
  87  // leaveAfterFailedEnter is identical to Leave, but is marked noinline to
  88  // prevent it from being inlined into Enter, since as of this writing inlining
  89  // Leave into Enter prevents Enter from being inlined into its callers.
  90  //
  91  //go:noinline
  92  func (g *Gate) leaveAfterFailedEnter() {
  93  	if atomic.AddInt32(&g.userCount, -1) == math.MinInt32 {
  94  		g.leaveClosed()
  95  	}
  96  }
  97  
  98  // Leave leaves the gate. This must only be called after a successful call to
  99  // Enter(). If the gate has been closed and this is the last one inside the
 100  // gate, it will notify the closer that the gate is done.
 101  //
 102  // This function is thread-safe.
 103  func (g *Gate) Leave() {
 104  	if atomic.AddInt32(&g.userCount, -1) == math.MinInt32 {
 105  		g.leaveClosed()
 106  	}
 107  }
 108  
 109  func (g *Gate) leaveClosed() {
 110  	if atomic.LoadUintptr(&g.closingG) == 0 {
 111  		return
 112  	}
 113  	if g := atomic.SwapUintptr(&g.closingG, 0); g > preparingG {
 114  		goready(g, 0)
 115  	}
 116  }
 117  
 118  // Close closes the gate, causing future calls to Enter to fail, and waits
 119  // until all goroutines that are currently inside the gate leave before
 120  // returning.
 121  //
 122  // Only one goroutine can call this function.
 123  func (g *Gate) Close() {
 124  	if atomic.LoadInt32(&g.userCount) == math.MinInt32 {
 125  		// The gate is already closed, with no goroutines inside. For legacy
 126  		// reasons, we have to allow Close to be called again in this case.
 127  		return
 128  	}
 129  	if v := atomic.AddInt32(&g.userCount, math.MinInt32); v == math.MinInt32 {
 130  		// userCount was already 0.
 131  		return
 132  	} else if v >= 0 {
 133  		panic("concurrent Close of sync.Gate")
 134  	}
 135  
 136  	if g := atomic.SwapUintptr(&g.closingG, preparingG); g != 0 {
 137  		panic(fmt.Sprintf("invalid sync.Gate.closingG during Close: %#x", g))
 138  	}
 139  	if atomic.LoadInt32(&g.userCount) == math.MinInt32 {
 140  		// The last call to Leave arrived while we were setting up closingG.
 141  		return
 142  	}
 143  	// WaitReasonSemacquire/TraceBlockSync are consistent with WaitGroup.
 144  	gopark(gateCommit, gohacks.Noescape(unsafe.Pointer(&g.closingG)), WaitReasonSemacquire, TraceBlockSync, 0)
 145  }
 146  
 147  //go:norace
 148  //go:nosplit
 149  func gateCommit(g uintptr, closingG unsafe.Pointer) bool {
 150  	return RaceUncheckedAtomicCompareAndSwapUintptr((*uintptr)(closingG), preparingG, g)
 151  }
 152