gopkg.in/rethinkdb/rethinkdb-go.v6@v6.2.2/internal/utils/conn_counter.go (about)

     1  package utils
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  	"regexp"
     8  	"runtime"
     9  	"runtime/debug"
    10  	"strings"
    11  	"sync"
    12  	"sync/atomic"
    13  	"time"
    14  )
    15  
    16  var connsCount int64
    17  var printer sync.Once
    18  
    19  var mu sync.Mutex
    20  var createdStacks = make(map[string]int)
    21  var reg = regexp.MustCompile(`0[xX][0-9a-fA-F]+`)
    22  
    23  type connCounting struct {
    24  	net.Conn
    25  
    26  	closed bool
    27  }
    28  
    29  // Socket leak debug net.Conn wrapper
    30  func NewCountingConn(conn net.Conn) net.Conn {
    31  	c := &connCounting{
    32  		Conn:   conn,
    33  		closed: false,
    34  	}
    35  	runtime.SetFinalizer(c, func(cc *connCounting) {
    36  		if !cc.closed {
    37  			atomic.AddInt64(&connsCount, -1)
    38  			cc.closed = true
    39  		}
    40  	})
    41  
    42  	atomic.AddInt64(&connsCount, 1)
    43  	printer.Do(func() {
    44  		go func() {
    45  			t := time.NewTicker(time.Second)
    46  			f, err := os.Create("sockets.ticker")
    47  			if err != nil {
    48  				fmt.Fprintf(os.Stderr, "Failed to create sockets.ticker file: %v\n", err)
    49  				return
    50  			}
    51  			for {
    52  				<-t.C
    53  				fmt.Fprintf(f, "Connections count: %v\n", atomic.LoadInt64(&connsCount))
    54  				f.Sync()
    55  			}
    56  		}()
    57  	})
    58  
    59  	st := string(debug.Stack())
    60  	st = st[strings.Index(st, "\n")+1:]
    61  	st = reg.ReplaceAllString(st, "")
    62  
    63  	mu.Lock()
    64  	_, has := createdStacks[st]
    65  	if !has {
    66  		createdStacks[st] = 1
    67  	} else {
    68  		createdStacks[st]++
    69  	}
    70  	printStacks()
    71  	mu.Unlock()
    72  
    73  	return c
    74  }
    75  
    76  func (c *connCounting) Close() error {
    77  	if !c.closed {
    78  		atomic.AddInt64(&connsCount, -1)
    79  		c.closed = true
    80  	}
    81  	return c.Conn.Close()
    82  }
    83  
    84  func printStacks() {
    85  	f, _ := os.Create("sockets.created")
    86  	for s, c := range createdStacks {
    87  		fmt.Fprintf(f, "%v:\n%v\n\n", c, s)
    88  	}
    89  	f.Sync()
    90  	f.Close()
    91  }