1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166
|
// Copyright 2016 Michal Witkowski. All Rights Reserved.
// See LICENSE for licensing terms.
package conntrack
import (
"context"
"fmt"
"net"
"sync"
"golang.org/x/net/trace"
)
var (
dialerNameKey = "conntrackDialerKey"
)
type dialerOpts struct {
name string
monitoring bool
tracing bool
parentDialContextFunc dialerContextFunc
}
type dialerOpt func(*dialerOpts)
type dialerContextFunc func(context.Context, string, string) (net.Conn, error)
// DialWithName sets the name of the dialer for tracking and monitoring.
// This is the name for the dialer (default is `default`), but for `NewDialContextFunc` can be overwritten from the
// Context using `DialNameToContext`.
func DialWithName(name string) dialerOpt {
return func(opts *dialerOpts) {
opts.name = name
}
}
// DialWithoutMonitoring turns *off* Prometheus monitoring for this dialer.
func DialWithoutMonitoring() dialerOpt {
return func(opts *dialerOpts) {
opts.monitoring = false
}
}
// DialWithTracing turns *on* the /debug/events tracing of the dial calls.
func DialWithTracing() dialerOpt {
return func(opts *dialerOpts) {
opts.tracing = true
}
}
// DialWithDialer allows you to override the `net.Dialer` instance used to actually conduct the dials.
func DialWithDialer(parentDialer *net.Dialer) dialerOpt {
return DialWithDialContextFunc(parentDialer.DialContext)
}
// DialWithDialContextFunc allows you to override func gets used for the actual dialing. The default is `net.Dialer.DialContext`.
func DialWithDialContextFunc(parentDialerFunc dialerContextFunc) dialerOpt {
return func(opts *dialerOpts) {
opts.parentDialContextFunc = parentDialerFunc
}
}
// DialNameFromContext returns the name of the dialer from the context of the DialContext func, if any.
func DialNameFromContext(ctx context.Context) string {
val, ok := ctx.Value(dialerNameKey).(string)
if !ok {
return ""
}
return val
}
// DialNameToContext returns a context that will contain a dialer name override.
func DialNameToContext(ctx context.Context, dialerName string) context.Context {
return context.WithValue(ctx, dialerNameKey, dialerName)
}
// NewDialContextFunc returns a `DialContext` function that tracks outbound connections.
// The signature is compatible with `http.Tranport.DialContext` and is meant to be used there.
func NewDialContextFunc(optFuncs ...dialerOpt) func(context.Context, string, string) (net.Conn, error) {
opts := &dialerOpts{name: defaultName, monitoring: true, parentDialContextFunc: (&net.Dialer{}).DialContext}
for _, f := range optFuncs {
f(opts)
}
if opts.monitoring {
PreRegisterDialerMetrics(opts.name)
}
return func(ctx context.Context, network string, addr string) (net.Conn, error) {
name := opts.name
if ctxName := DialNameFromContext(ctx); ctxName != "" {
name = ctxName
}
return dialClientConnTracker(ctx, network, addr, name, opts)
}
}
// NewDialFunc returns a `Dial` function that tracks outbound connections.
// The signature is compatible with `http.Tranport.Dial` and is meant to be used there for Go < 1.7.
func NewDialFunc(optFuncs ...dialerOpt) func(string, string) (net.Conn, error) {
dialContextFunc := NewDialContextFunc(optFuncs...)
return func(network string, addr string) (net.Conn, error) {
return dialContextFunc(context.TODO(), network, addr)
}
}
type clientConnTracker struct {
net.Conn
opts *dialerOpts
dialerName string
event trace.EventLog
mu sync.Mutex
}
func dialClientConnTracker(ctx context.Context, network string, addr string, dialerName string, opts *dialerOpts) (net.Conn, error) {
var event trace.EventLog
if opts.tracing {
event = trace.NewEventLog(fmt.Sprintf("net.ClientConn.%s", dialerName), fmt.Sprintf("%v", addr))
}
if opts.monitoring {
reportDialerConnAttempt(dialerName)
}
conn, err := opts.parentDialContextFunc(ctx, network, addr)
if err != nil {
if event != nil {
event.Errorf("failed dialing: %v", err)
event.Finish()
}
if opts.monitoring {
reportDialerConnFailed(dialerName, err)
}
return nil, err
}
if event != nil {
event.Printf("established: %s -> %s", conn.LocalAddr(), conn.RemoteAddr())
}
if opts.monitoring {
reportDialerConnEstablished(dialerName)
}
tracker := &clientConnTracker{
Conn: conn,
opts: opts,
dialerName: dialerName,
event: event,
}
return tracker, nil
}
func (ct *clientConnTracker) Close() error {
err := ct.Conn.Close()
ct.mu.Lock()
if ct.event != nil {
if err != nil {
ct.event.Errorf("failed closing: %v", err)
} else {
ct.event.Printf("closing")
}
ct.event.Finish()
ct.event = nil
}
ct.mu.Unlock()
if ct.opts.monitoring {
reportDialerConnClosed(ct.dialerName)
}
return err
}
|