github.com/go-asm/go@v1.21.1-0.20240213172139-40c5ead50c48/cmd/go/base/limit.go (about)

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