github.com/unigraph-dev/dgraph@v1.1.1-0.20200923154953-8b52b426f765/xidmap/xidmap_test.go (about)

     1  package xidmap
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  	"os"
     7  	"runtime"
     8  	"sync"
     9  	"sync/atomic"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/dgraph-io/badger"
    14  	"github.com/dgraph-io/dgraph/testutil"
    15  	"github.com/dgraph-io/dgraph/x"
    16  	"github.com/stretchr/testify/require"
    17  )
    18  
    19  // Opens a badger db and runs a a test on it.
    20  func withDB(t *testing.T, test func(db *badger.DB)) {
    21  	dir, err := ioutil.TempDir(".", "badger-test")
    22  	require.NoError(t, err)
    23  	defer os.RemoveAll(dir)
    24  
    25  	opt := badger.LSMOnlyOptions(dir)
    26  	db, err := badger.Open(opt)
    27  	require.NoError(t, err)
    28  	defer db.Close()
    29  
    30  	test(db)
    31  }
    32  
    33  func TestXidmap(t *testing.T) {
    34  	conn, err := x.SetupConnection(testutil.SockAddrZero, nil, false)
    35  	require.NoError(t, err)
    36  	require.NotNil(t, conn)
    37  
    38  	withDB(t, func(db *badger.DB) {
    39  		xidmap := New(conn, db)
    40  
    41  		uida := xidmap.AssignUid("a")
    42  		require.Equal(t, uida, xidmap.AssignUid("a"))
    43  
    44  		uidb := xidmap.AssignUid("b")
    45  		require.True(t, uida != uidb)
    46  		require.Equal(t, uidb, xidmap.AssignUid("b"))
    47  
    48  		to := xidmap.AllocateUid() + uint64(1e6+3)
    49  		xidmap.BumpTo(to)
    50  		uid := xidmap.AllocateUid() // Does not have to be above the bump.
    51  		t.Logf("bump up to: %d. allocated: %d", to, uid)
    52  
    53  		require.NoError(t, xidmap.Flush())
    54  		xidmap = nil
    55  
    56  		xidmap2 := New(conn, db)
    57  		require.Equal(t, uida, xidmap2.AssignUid("a"))
    58  		require.Equal(t, uidb, xidmap2.AssignUid("b"))
    59  	})
    60  }
    61  
    62  func TestXidmapMemory(t *testing.T) {
    63  	if testing.Short() {
    64  		t.Skip("skipping because -short=true")
    65  	}
    66  
    67  	var loop uint32
    68  	bToMb := func(b uint64) uint64 {
    69  		return b / 1024 / 1024
    70  	}
    71  	printMemory := func() {
    72  		var m runtime.MemStats
    73  		runtime.ReadMemStats(&m)
    74  		// For info on each, see: https://golang.org/pkg/runtime/#MemStats
    75  		fmt.Printf(" Heap = %v M", bToMb(m.HeapInuse))
    76  		fmt.Printf(" Alloc = %v M", bToMb(m.Alloc))
    77  		fmt.Printf(" Sys = %v M", bToMb(m.Sys))
    78  		fmt.Printf(" Loop = %.2fM", float64(atomic.LoadUint32(&loop))/1e6)
    79  		fmt.Printf(" NumGC = %v\n", m.NumGC)
    80  	}
    81  	go func() {
    82  		ticker := time.NewTicker(time.Second)
    83  		defer ticker.Stop()
    84  
    85  		for range ticker.C {
    86  			printMemory()
    87  		}
    88  	}()
    89  
    90  	conn, err := x.SetupConnection(testutil.SockAddrZero, nil, false)
    91  	require.NoError(t, err)
    92  	require.NotNil(t, conn)
    93  
    94  	xidmap := New(conn, nil)
    95  
    96  	start := time.Now()
    97  	var wg sync.WaitGroup
    98  	for numGo := 0; numGo < 32; numGo++ {
    99  		wg.Add(1)
   100  		go func() {
   101  			defer wg.Done()
   102  			for {
   103  				i := atomic.AddUint32(&loop, 1)
   104  				if i > 50e6 {
   105  					return
   106  				}
   107  				xidmap.AssignUid(fmt.Sprintf("xid-%d", i))
   108  			}
   109  		}()
   110  	}
   111  	wg.Wait()
   112  	t.Logf("Time taken: %v", time.Since(start).Round(time.Millisecond))
   113  }
   114  
   115  func BenchmarkXidmap(b *testing.B) {
   116  	conn, err := x.SetupConnection(testutil.SockAddrZero, nil, false)
   117  	x.Check(err)
   118  	x.AssertTrue(conn != nil)
   119  
   120  	var counter uint64
   121  	xidmap := New(conn, nil)
   122  	b.ResetTimer()
   123  
   124  	b.RunParallel(func(pb *testing.PB) {
   125  		for pb.Next() {
   126  			xid := atomic.AddUint64(&counter, 1)
   127  			xidmap.AssignUid(fmt.Sprintf("xid-%d", xid))
   128  		}
   129  	})
   130  }