limit.go

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package base
     6  
     7  import (
     8  	"fmt"
     9  	"internal/godebug"
    10  	"runtime"
    11  	"strconv"
    12  	"sync"
    13  )
    14  
    15  var NetLimitGodebug = godebug.New("#cmdgonetlimit")
    16  
    17  // NetLimit returns the limit on concurrent network operations
    18  // configured by GODEBUG=cmdgonetlimit, if any.
    19  //
    20  // A limit of 0 (indicated by 0, true) means that network operations should not
    21  // be allowed.
    22  func NetLimit() (int, bool) {
    23  	netLimitOnce.Do(func() {
    24  		s := NetLimitGodebug.Value()
    25  		if s == "" {
    26  			return
    27  		}
    28  
    29  		n, err := strconv.Atoi(s)
    30  		if err != nil {
    31  			Fatalf("invalid %s: %v", NetLimitGodebug.Name(), err)
    32  		}
    33  		if n < 0 {
    34  			// Treat negative values as unlimited.
    35  			return
    36  		}
    37  		netLimitSem = make(chan struct{}, n)
    38  	})
    39  
    40  	return cap(netLimitSem), netLimitSem != nil
    41  }
    42  
    43  // AcquireNet acquires a semaphore token for a network operation.
    44  func AcquireNet() (release func(), err error) {
    45  	hasToken := false
    46  	if n, ok := NetLimit(); ok {
    47  		if n == 0 {
    48  			return nil, fmt.Errorf("network disabled by %v=%v", NetLimitGodebug.Name(), NetLimitGodebug.Value())
    49  		}
    50  		netLimitSem <- struct{}{}
    51  		hasToken = true
    52  	}
    53  
    54  	checker := new(netTokenChecker)
    55  	runtime.SetFinalizer(checker, (*netTokenChecker).panicUnreleased)
    56  
    57  	return func() {
    58  		if checker.released {
    59  			panic("internal error: net token released twice")
    60  		}
    61  		checker.released = true
    62  		if hasToken {
    63  			<-netLimitSem
    64  		}
    65  		runtime.SetFinalizer(checker, nil)
    66  	}, nil
    67  }
    68  
    69  var (
    70  	netLimitOnce sync.Once
    71  	netLimitSem  chan struct{}
    72  )
    73  
    74  type netTokenChecker struct {
    75  	released bool
    76  	// We want to use a finalizer to check that all acquired tokens are returned,
    77  	// so we arbitrarily pad the tokens with a string to defeat the runtime's
    78  	// “tiny allocator”.
    79  	unusedAvoidTinyAllocator string
    80  }
    81  
    82  func (c *netTokenChecker) panicUnreleased() {
    83  	panic("internal error: net token acquired but not released")
    84  }
    85
View as plain text