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
5package base
6
7import (
8	"fmt"
9	"internal/godebug"
10	"runtime"
11	"strconv"
12	"sync"
13)
14
15var 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.
22func 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.
44func 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
69var (
70	netLimitOnce sync.Once
71	netLimitSem  chan struct{}
72)
73
74type 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
82func (c *netTokenChecker) panicUnreleased() {
83	panic("internal error: net token acquired but not released")
84}
85