src/cmd/go/internal/base/limit.go

// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package base

import (
	"fmt"
	"internal/godebug"
	"runtime"
	"strconv"
	"sync"
)

var NetLimitGodebug = godebug.New("#cmdgonetlimit")

// NetLimit returns the limit on concurrent network operations
// configured by GODEBUG=cmdgonetlimit, if any.
//
// A limit of 0 (indicated by 0, true) means that network operations should not
// be allowed.
func NetLimit() (int, bool) {
	netLimitOnce.Do(func() {
		s := NetLimitGodebug.Value()
		if s == "" {
			return
		}

		n, err := strconv.Atoi(s)
		if err != nil {
			Fatalf("invalid %s: %v", NetLimitGodebug.Name(), err)
		}
		if n < 0 {
			// Treat negative values as unlimited.
			return
		}
		netLimitSem = make(chan struct{}, n)
	})

	return cap(netLimitSem), netLimitSem != nil
}

// AcquireNet acquires a semaphore token for a network operation.
func AcquireNet() (release func(), err error) {
	hasToken := false
	if n, ok := NetLimit(); ok {
		if n == 0 {
			return nil, fmt.Errorf("network disabled by %v=%v", NetLimitGodebug.Name(), NetLimitGodebug.Value())
		}
		netLimitSem <- struct{}{}
		hasToken = true
	}

	checker := new(netTokenChecker)
	cleanup := runtime.AddCleanup(checker, func(_ int) { panic("internal error: net token acquired but not released") }, 0)

	return func() {
		if checker.released {
			panic("internal error: net token released twice")
		}
		checker.released = true
		if hasToken {
			<-netLimitSem
		}
		cleanup.Stop()

		// checker may be dead at the moment after we last access
		// it in this function, so the cleanup can fire before Stop
		// completes. Keep checker alive while we call Stop. See
		// the documentation for runtime.Cleanup.Stop.
		runtime.KeepAlive(checker)
	}, nil
}

var (
	netLimitOnce sync.Once
	netLimitSem  chan struct{}
)

type netTokenChecker struct {
	released bool
	// We want to use a finalizer to check that all acquired tokens are returned,
	// so we arbitrarily pad the tokens with a string to defeat the runtime's
	// “tiny allocator”.
	unusedAvoidTinyAllocator string
}