github.com/fjl/memsize@v0.0.2/memsize.go (about)

     1  package memsize
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"reflect"
     7  	"sort"
     8  	"strings"
     9  	"text/tabwriter"
    10  	"unsafe"
    11  )
    12  
    13  // Scan traverses all objects reachable from v and counts how much memory
    14  // is used per type. The value must be a non-nil pointer to any value.
    15  func Scan(v interface{}) Sizes {
    16  	rv := reflect.ValueOf(v)
    17  	if rv.Kind() != reflect.Ptr || rv.IsNil() {
    18  		panic("value to scan must be non-nil pointer")
    19  	}
    20  
    21  	stopTheWorld(stwReadMemStats)
    22  	defer startTheWorld()
    23  
    24  	ctx := newContext()
    25  	ctx.scan(invalidAddr, rv, false)
    26  	ctx.s.BitmapSize = ctx.seen.size()
    27  	ctx.s.BitmapUtilization = ctx.seen.utilization()
    28  	return *ctx.s
    29  }
    30  
    31  // Sizes is the result of a scan.
    32  type Sizes struct {
    33  	Total  uintptr
    34  	ByType map[reflect.Type]*TypeSize
    35  	// Internal stats (for debugging)
    36  	BitmapSize        uintptr
    37  	BitmapUtilization float32
    38  }
    39  
    40  type TypeSize struct {
    41  	Total uintptr
    42  	Count uintptr
    43  }
    44  
    45  func newSizes() *Sizes {
    46  	return &Sizes{ByType: make(map[reflect.Type]*TypeSize)}
    47  }
    48  
    49  // Report returns a human-readable report.
    50  func (s Sizes) Report() string {
    51  	type typLine struct {
    52  		name  string
    53  		count uintptr
    54  		total uintptr
    55  	}
    56  	tab := []typLine{{"ALL", 0, s.Total}}
    57  	for _, typ := range s.ByType {
    58  		tab[0].count += typ.Count
    59  	}
    60  	maxname := 0
    61  	for typ, s := range s.ByType {
    62  		line := typLine{typ.String(), s.Count, s.Total}
    63  		tab = append(tab, line)
    64  		if len(line.name) > maxname {
    65  			maxname = len(line.name)
    66  		}
    67  	}
    68  	sort.Slice(tab, func(i, j int) bool { return tab[i].total > tab[j].total })
    69  
    70  	buf := new(bytes.Buffer)
    71  	w := tabwriter.NewWriter(buf, 0, 0, 0, ' ', tabwriter.AlignRight)
    72  	for _, line := range tab {
    73  		namespace := strings.Repeat(" ", maxname-len(line.name))
    74  		fmt.Fprintf(w, "%s%s\t  %v\t  %s\t\n", line.name, namespace, line.count, HumanSize(line.total))
    75  	}
    76  	w.Flush()
    77  	return buf.String()
    78  }
    79  
    80  // addValue is called during scan and adds the memory of given object.
    81  func (s *Sizes) addValue(v reflect.Value, size uintptr) {
    82  	s.Total += size
    83  	rs := s.ByType[v.Type()]
    84  	if rs == nil {
    85  		rs = new(TypeSize)
    86  		s.ByType[v.Type()] = rs
    87  	}
    88  	rs.Total += size
    89  	rs.Count++
    90  }
    91  
    92  type context struct {
    93  	// We track previously scanned objects to prevent infinite loops
    94  	// when scanning cycles and to prevent counting objects more than once.
    95  	seen *bitmap
    96  	tc   typCache
    97  	s    *Sizes
    98  }
    99  
   100  func newContext() *context {
   101  	return &context{seen: newBitmap(), tc: make(typCache), s: newSizes()}
   102  }
   103  
   104  // scan walks all objects below v, determining their size. It returns the size of the
   105  // previously unscanned parts of the object.
   106  func (c *context) scan(addr address, v reflect.Value, add bool) (extraSize uintptr) {
   107  	size := v.Type().Size()
   108  	var marked uintptr
   109  	if addr.valid() {
   110  		marked = c.seen.countRange(uintptr(addr), size)
   111  		if marked == size {
   112  			return 0 // Skip if we have already seen the whole object.
   113  		}
   114  		c.seen.markRange(uintptr(addr), size)
   115  	}
   116  	// fmt.Printf("%v: %v ⮑ (marked %d)\n", addr, v.Type(), marked)
   117  	if c.tc.needScan(v.Type()) {
   118  		extraSize = c.scanContent(addr, v)
   119  	}
   120  	size -= marked
   121  	size += extraSize
   122  	// fmt.Printf("%v: %v %d (add %v, size %d, marked %d, extra %d)\n", addr, v.Type(), size+extraSize, add, v.Type().Size(), marked, extraSize)
   123  	if add {
   124  		c.s.addValue(v, size)
   125  	}
   126  	return size
   127  }
   128  
   129  // scanContent and all other scan* functions below return the amount of 'extra' memory
   130  // (e.g. slice data) that is referenced by the object.
   131  func (c *context) scanContent(addr address, v reflect.Value) uintptr {
   132  	switch v.Kind() {
   133  	case reflect.Array:
   134  		return c.scanArray(addr, v)
   135  	case reflect.Chan:
   136  		return c.scanChan(v)
   137  	case reflect.Func:
   138  		// can't do anything here
   139  		return 0
   140  	case reflect.Interface:
   141  		return c.scanInterface(v)
   142  	case reflect.Map:
   143  		return c.scanMap(v)
   144  	case reflect.Ptr:
   145  		if !v.IsNil() {
   146  			c.scan(address(v.Pointer()), v.Elem(), true)
   147  		}
   148  		return 0
   149  	case reflect.Slice:
   150  		return c.scanSlice(v)
   151  	case reflect.String:
   152  		return uintptr(v.Len())
   153  	case reflect.Struct:
   154  		return c.scanStruct(addr, v)
   155  	default:
   156  		unhandledKind(v.Kind())
   157  		return 0
   158  	}
   159  }
   160  
   161  func (c *context) scanChan(v reflect.Value) uintptr {
   162  	etyp := v.Type().Elem()
   163  	extra := uintptr(0)
   164  	if c.tc.needScan(etyp) {
   165  		// Scan the channel buffer. This is unsafe but doesn't race because
   166  		// the world is stopped during scan.
   167  		hchan := unsafe.Pointer(v.Pointer())
   168  		for i := uint(0); i < uint(v.Cap()); i++ {
   169  			addr := chanbuf(hchan, i)
   170  			elem := reflect.NewAt(etyp, addr).Elem()
   171  			extra += c.scanContent(address(addr), elem)
   172  		}
   173  	}
   174  	return uintptr(v.Cap())*etyp.Size() + extra
   175  }
   176  
   177  func (c *context) scanStruct(base address, v reflect.Value) uintptr {
   178  	extra := uintptr(0)
   179  	for i := 0; i < v.NumField(); i++ {
   180  		f := v.Type().Field(i)
   181  		if c.tc.needScan(f.Type) {
   182  			addr := base.addOffset(f.Offset)
   183  			extra += c.scanContent(addr, v.Field(i))
   184  		}
   185  	}
   186  	return extra
   187  }
   188  
   189  func (c *context) scanArray(addr address, v reflect.Value) uintptr {
   190  	esize := v.Type().Elem().Size()
   191  	extra := uintptr(0)
   192  	for i := 0; i < v.Len(); i++ {
   193  		extra += c.scanContent(addr, v.Index(i))
   194  		addr = addr.addOffset(esize)
   195  	}
   196  	return extra
   197  }
   198  
   199  func (c *context) scanSlice(v reflect.Value) uintptr {
   200  	slice := v.Slice(0, v.Cap())
   201  	esize := slice.Type().Elem().Size()
   202  	base := slice.Pointer()
   203  	// Add size of the unscanned portion of the backing array to extra.
   204  	blen := uintptr(slice.Len()) * esize
   205  	marked := c.seen.countRange(base, blen)
   206  	extra := blen - marked
   207  	c.seen.markRange(uintptr(base), blen)
   208  	if c.tc.needScan(slice.Type().Elem()) {
   209  		// Elements may contain pointers, scan them individually.
   210  		addr := address(base)
   211  		for i := 0; i < slice.Len(); i++ {
   212  			extra += c.scanContent(addr, slice.Index(i))
   213  			addr = addr.addOffset(esize)
   214  		}
   215  	}
   216  	return extra
   217  }
   218  
   219  func (c *context) scanMap(v reflect.Value) uintptr {
   220  	var (
   221  		typ   = v.Type()
   222  		len   = uintptr(v.Len())
   223  		extra = uintptr(0)
   224  	)
   225  	if c.tc.needScan(typ.Key()) || c.tc.needScan(typ.Elem()) {
   226  		iterateMap(v, func(k, v reflect.Value) {
   227  			extra += c.scan(invalidAddr, k, false)
   228  			extra += c.scan(invalidAddr, v, false)
   229  		})
   230  	} else {
   231  		extra = len*typ.Key().Size() + len*typ.Elem().Size()
   232  	}
   233  	return extra
   234  }
   235  
   236  func (c *context) scanInterface(v reflect.Value) uintptr {
   237  	elem := v.Elem()
   238  	if !elem.IsValid() {
   239  		return 0 // nil interface
   240  	}
   241  	extra := c.scan(invalidAddr, elem, false)
   242  	if elem.Type().Kind() == reflect.Ptr {
   243  		extra -= uintptrBytes
   244  	}
   245  	return extra
   246  }