github.com/dgraph-io/ristretto@v0.1.2-0.20240116140435-c67e07994f91/contrib/memtest/main.go (about)

     1  /*
     2   * Copyright 2020 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package main
    18  
    19  // #include <stdlib.h>
    20  import "C"
    21  import (
    22  	"fmt"
    23  	"log"
    24  	"math/rand"
    25  	"net/http"
    26  	_ "net/http/pprof"
    27  	"os"
    28  	"os/signal"
    29  	"runtime"
    30  	"sync/atomic"
    31  	"syscall"
    32  	"time"
    33  	"unsafe"
    34  
    35  	"github.com/dustin/go-humanize"
    36  
    37  	"github.com/dgraph-io/ristretto/z"
    38  )
    39  
    40  type S struct {
    41  	key  uint64
    42  	val  []byte
    43  	next *S
    44  	inGo bool
    45  }
    46  
    47  var (
    48  	ssz      = int(unsafe.Sizeof(S{}))
    49  	lo, hi   = int64(1 << 30), int64(16 << 30)
    50  	increase = true
    51  	stop     int32
    52  	fill     []byte
    53  	maxMB    = 32
    54  
    55  	cycles int64 = 16
    56  )
    57  var numbytes int64
    58  var counter int64
    59  
    60  func newS(sz int) *S {
    61  	var s *S
    62  	if b := Calloc(ssz); len(b) > 0 {
    63  		s = (*S)(unsafe.Pointer(&b[0]))
    64  	} else {
    65  		s = &S{inGo: true}
    66  	}
    67  
    68  	s.val = Calloc(sz)
    69  	copy(s.val, fill)
    70  	if s.next != nil {
    71  		log.Fatalf("news.next must be nil: %p", s.next)
    72  	}
    73  	return s
    74  }
    75  
    76  func freeS(s *S) {
    77  	Free(s.val)
    78  	if !s.inGo {
    79  		buf := (*[z.MaxArrayLen]byte)(unsafe.Pointer(s))[:ssz:ssz]
    80  		Free(buf)
    81  	}
    82  }
    83  
    84  func (s *S) allocateNext(sz int) {
    85  	ns := newS(sz)
    86  	s.next, ns.next = ns, s.next
    87  }
    88  
    89  func (s *S) deallocNext() {
    90  	if s.next == nil {
    91  		log.Fatal("next should not be nil")
    92  	}
    93  	next := s.next
    94  	s.next = next.next
    95  	freeS(next)
    96  }
    97  
    98  func memory() {
    99  	// In normal mode, z.NumAllocBytes would always be zero. So, this program would misbehave.
   100  	curMem := NumAllocBytes()
   101  	if increase {
   102  		if curMem > hi {
   103  			increase = false
   104  		}
   105  	} else {
   106  		if curMem < lo {
   107  			increase = true
   108  			runtime.GC()
   109  			time.Sleep(3 * time.Second)
   110  
   111  			counter++
   112  		}
   113  	}
   114  	var js z.MemStats
   115  	z.ReadMemStats(&js)
   116  
   117  	fmt.Printf("[%d] Current Memory: %s. Increase? %v, MemStats [Active: %s, Allocated: %s,"+
   118  		" Resident: %s, Retained: %s]\n",
   119  		counter, humanize.IBytes(uint64(curMem)), increase,
   120  		humanize.IBytes(js.Active), humanize.IBytes(js.Allocated),
   121  		humanize.IBytes(js.Resident), humanize.IBytes(js.Retained))
   122  }
   123  
   124  func viaLL() {
   125  	ticker := time.NewTicker(10 * time.Millisecond)
   126  	defer ticker.Stop()
   127  
   128  	root := newS(1)
   129  	for range ticker.C {
   130  		if counter >= cycles {
   131  			fmt.Printf("Finished %d cycles. Deallocating...\n", counter)
   132  			break
   133  		}
   134  		if atomic.LoadInt32(&stop) == 1 {
   135  			break
   136  		}
   137  		if increase {
   138  			root.allocateNext(rand.Intn(maxMB) << 20)
   139  		} else {
   140  			root.deallocNext()
   141  		}
   142  		memory()
   143  	}
   144  	for root.next != nil {
   145  		root.deallocNext()
   146  		memory()
   147  	}
   148  	freeS(root)
   149  }
   150  
   151  func main() {
   152  	check()
   153  	fill = make([]byte, maxMB<<20)
   154  	_, _ = rand.Read(fill)
   155  
   156  	c := make(chan os.Signal, 10)
   157  	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
   158  	go func() {
   159  		<-c
   160  		fmt.Println("Stopping")
   161  		atomic.StoreInt32(&stop, 1)
   162  	}()
   163  	go func() {
   164  		if err := http.ListenAndServe("0.0.0.0:8080", nil); err != nil {
   165  			log.Fatalf("Error: %v", err)
   166  		}
   167  	}()
   168  
   169  	viaLL()
   170  	if left := NumAllocBytes(); left != 0 {
   171  		log.Fatalf("Unable to deallocate all memory: %v\n", left)
   172  	}
   173  	runtime.GC()
   174  	fmt.Println("Done. Reduced to zero memory usage.")
   175  	time.Sleep(5 * time.Second)
   176  }