Golang store.Store類代碼示例

func main() {
	flag.Parse()
	if !flag.Parsed() {
		glog.Fatal("Unable to parse flags")
	}
	logrus.SetLevel(logrus.InfoLevel)
	numCpus := runtime.NumCPU()
	prev := runtime.GOMAXPROCS(numCpus)
	glog.WithField("num_cpu", numCpus).
		WithField("prev_maxprocs", prev).
		Info("Set max procs to num cpus")

	ps := new(store.Store)
	ps.Init(*postingDir)
	defer ps.Close()

	clog := commit.NewLogger(*mutationDir, "dgraph", 50<<20)
	clog.SyncEvery = 1
	clog.Init()
	defer clog.Close()

	posting.Init(ps, clog)

	http.HandleFunc("/query", queryHandler)
	glog.WithField("port", *port).Info("Listening for requests...")
	if err := http.ListenAndServe(":"+*port, nil); err != nil {
		x.Err(glog, err).Fatal("ListenAndServe")
	}
}

func prepare() (dir1, dir2 string, clog *commit.Logger, rerr error) {
	var err error
	dir1, err = ioutil.TempDir("", "storetest_")
	if err != nil {
		return "", "", nil, err
	}
	ps := new(store.Store)
	ps.Init(dir1)

	dir2, err = ioutil.TempDir("", "storemuts_")
	if err != nil {
		return dir1, "", nil, err
	}
	clog = commit.NewLogger(dir2, "mutations", 50<<20)
	clog.Init()

	posting.Init(ps, clog)

	f, err := os.Open("testdata.nq")
	if err != nil {
		return dir1, dir2, clog, err
	}
	defer f.Close()
	_, err = loader.HandleRdfReader(f, 0, 1)
	if err != nil {
		return dir1, dir2, clog, err
	}
	return dir1, dir2, clog, nil
}

func benchmarkAddMutations(n int, b *testing.B) {
	// logrus.SetLevel(logrus.DebugLevel)
	l := NewList()
	key := Key(1, "name")
	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		b.Error(err)
		return
	}

	defer os.RemoveAll(dir)
	ps := new(store.Store)
	ps.Init(dir)

	clog := commit.NewLogger(dir, "mutations", 50<<20)
	clog.SyncEvery = n
	clog.Init()
	defer clog.Close()
	l.init(key, ps, clog)
	b.ResetTimer()

	ts := time.Now()
	for i := 0; i < b.N; i++ {
		edge := x.DirectedEdge{
			ValueId:   uint64(rand.Intn(b.N) + 1),
			Source:    "testing",
			Timestamp: ts.Add(time.Microsecond),
		}
		if err := l.AddMutation(edge, Set); err != nil {
			b.Error(err)
		}
	}
}

func main() {
	flag.Parse()
	if !flag.Parsed() {
		glog.Fatal("Unable to parse flags")
	}
	if len(*cpuprofile) > 0 {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			glog.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	logrus.SetLevel(logrus.InfoLevel)
	numCpus := *numcpu
	prevProcs := runtime.GOMAXPROCS(numCpus)
	glog.WithField("num_cpus", numCpus).
		WithField("prev_maxprocs", prevProcs).
		Info("Set max procs to num  cpus")

	if len(*rdfGzips) == 0 {
		glog.Fatal("No RDF GZIP files specified")
	}

	ps := new(store.Store)
	ps.Init(*uidDir)
	defer ps.Close()

	posting.Init(ps, nil)

	files := strings.Split(*rdfGzips, ",")
	for _, path := range files {
		if len(path) == 0 {
			continue
		}
		glog.WithField("path", path).Info("Handling...")
		f, err := os.Open(path)
		if err != nil {
			glog.WithError(err).Fatal("Unable to open rdf file.")
		}

		r, err := gzip.NewReader(f)
		if err != nil {
			glog.WithError(err).Fatal("Unable to create gzip reader.")
		}

		count, err := loader.HandleRdfReaderWhileAssign(r, *mod)
		if err != nil {
			glog.WithError(err).Fatal("While handling rdf reader.")
		}
		glog.WithField("count", count).Info("RDFs parsed")
		r.Close()
		f.Close()
	}
	glog.Info("Calling merge lists")
	posting.MergeLists(100 * numCpus) // 100 per core.
}

func checkShard(ps *store.Store) (int, []byte) {
	it := ps.NewIterator()
	defer it.Close()

	count := 0
	for it.SeekToFirst(); it.Valid(); it.Next() {
		count++
	}
	return count, it.Key().Data()
}

func TestQuery(t *testing.T) {
	var numInstances uint64 = 2
	mod := math.MaxUint64 / numInstances
	minIdx0 := 0 * mod
	minIdx1 := 1 * mod

	logrus.SetLevel(logrus.DebugLevel)
	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return
	}
	defer os.RemoveAll(dir)

	ps := new(store.Store)
	ps.Init(dir)
	clog := commit.NewLogger(dir, "mutations", 50<<20)
	clog.Init()
	defer clog.Close()
	posting.Init(clog)

	uid.Init(ps)

	list := []string{"alice", "bob", "mallory", "ash", "man", "dgraph"}
	for _, str := range list {
		if farm.Fingerprint64([]byte(str))%numInstances == 0 {
			uid, err := rdf.GetUid(str, 0, numInstances)
			if uid < minIdx0 || uid > minIdx0+mod-1 {
				t.Error("Not the correct UID", err)
			}
			t.Logf("Instance-0 Correct UID", str, uid)

		} else {
			uid, err := rdf.GetUid(str, 1, numInstances)
			if uid < minIdx1 || uid > minIdx1+mod-1 {
				t.Error("Not the correct UID", err)
			}
			t.Logf("Instance-1 Correct UID", str, uid)
		}
	}
}

func TestNewGraph(t *testing.T) {
	var ex uint64
	ex = 101

	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return
	}

	ps := new(store.Store)
	ps.Init(dir)
	sg, err := newGraph(ex, "", ps)
	if err != nil {
		t.Error(err)
	}

	worker.Init(ps)

	uo := flatbuffers.GetUOffsetT(sg.result)
	r := new(task.Result)
	r.Init(sg.result, uo)
	if r.UidmatrixLength() != 1 {
		t.Errorf("Expected length 1. Got: %v", r.UidmatrixLength())
	}
	var ul task.UidList
	if ok := r.Uidmatrix(&ul, 0); !ok {
		t.Errorf("Unable to parse uidlist at index 0")
	}
	if ul.UidsLength() != 1 {
		t.Errorf("Expected length 1. Got: %v", ul.UidsLength())
	}
	if ul.Uids(0) != ex {
		t.Errorf("Expected uid: %v. Got: %v", ex, ul.Uids(0))
	}
}

func TestGetOrAssign(t *testing.T) {
	logrus.SetLevel(logrus.DebugLevel)

	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return
	}
	defer os.RemoveAll(dir)
	ps := new(store.Store)
	ps.Init(dir)
	clog := commit.NewLogger(dir, "mutations", 50<<20)
	clog.Init()
	defer clog.Close()

	posting.Init(ps, clog)

	var u1, u2 uint64
	{
		uid, err := GetOrAssign("externalid0")
		if err != nil {
			t.Error(err)
		}
		t.Logf("Found uid: [%x]", uid)
		u1 = uid
	}

	{
		uid, err := GetOrAssign("externalid1")
		if err != nil {
			t.Error(err)
		}
		t.Logf("Found uid: [%x]", uid)
		u2 = uid
	}

	if u1 == u2 {
		t.Error("Uid1 and Uid2 shouldn't be the same")
	}
	// return

	{
		uid, err := GetOrAssign("externalid0")
		if err != nil {
			t.Error(err)
		}
		t.Logf("Found uid: [%x]", uid)
		if u1 != uid {
			t.Error("Uid should be the same.")
		}
	}
	// return

	{
		xid, err := ExternalId(u1)
		if err != nil {
			t.Error(err)
		}
		if xid != "externalid0" {
			t.Errorf("Expected externalid0. Found: [%q]", xid)
		}
	}
	return
	{
		xid, err := ExternalId(u2)
		if err != nil {
			t.Error(err)
		}
		if xid != "externalid1" {
			t.Errorf("Expected externalid1. Found: [%q]", xid)
		}
	}
}

func TestAddMutation(t *testing.T) {
	// logrus.SetLevel(logrus.DebugLevel)
	l := NewList()
	key := Key(1, "name")
	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return
	}

	defer os.RemoveAll(dir)
	ps := new(store.Store)
	ps.Init(dir)

	clog := commit.NewLogger(dir, "mutations", 50<<20)
	clog.Init()
	defer clog.Close()

	l.init(key, ps, clog)

	edge := x.DirectedEdge{
		ValueId:   9,
		Source:    "testing",
		Timestamp: time.Now(),
	}
	if err := l.AddMutation(edge, Set); err != nil {
		t.Error(err)
	}
	/*
		if err := l.CommitIfDirty(); err != nil {
			t.Error(err)
		}
	*/

	if l.Length() != 1 {
		t.Error("Unable to find added elements in posting list")
	}
	var p types.Posting
	if ok := l.Get(&p, 0); !ok {
		t.Error("Unable to retrieve posting at 1st iter")
		t.Fail()
	}
	if p.Uid() != 9 {
		t.Errorf("Expected 9. Got: %v", p.Uid)
	}
	if string(p.Source()) != "testing" {
		t.Errorf("Expected testing. Got: %v", string(p.Source()))
	}
	// return // Test 1.

	// Add another edge now.
	edge.ValueId = 81
	l.AddMutation(edge, Set)
	// l.CommitIfDirty()
	if l.Length() != 2 {
		t.Errorf("Length: %d", l.Length())
		t.Fail()
	}

	var uid uint64
	uid = 1
	for i := 0; i < l.Length(); i++ {
		if ok := l.Get(&p, i); !ok {
			t.Error("Unable to retrieve posting at 2nd iter")
		}
		uid *= 9
		if p.Uid() != uid {
			t.Logf("Expected: %v. Got: %v", uid, p.Uid())
		}
	}
	// return // Test 2.

	// Add another edge, in between the two above.
	uids := []uint64{
		9, 49, 81,
	}
	edge.ValueId = 49
	if err := l.AddMutation(edge, Set); err != nil {
		t.Error(err)
	}
	/*
		if err := l.CommitIfDirty(); err != nil {
			t.Error(err)
		}
	*/
	if err := checkUids(t, l, uids...); err != nil {
		t.Error(err)
	}
	// return // Test 3.

	// Delete an edge, add an edge, replace an edge
	edge.ValueId = 49
	if err := l.AddMutation(edge, Del); err != nil {
		t.Error(err)
	}

	edge.ValueId = 69
	if err := l.AddMutation(edge, Set); err != nil {
		t.Error(err)
	}
//.........這裏部分代碼省略.........

func TestAddMutation_Value(t *testing.T) {
	// logrus.SetLevel(logrus.DebugLevel)
	glog.Debug("Running init...")
	ol := NewList()
	key := Key(10, "value")
	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return
	}

	defer os.RemoveAll(dir)
	ps := new(store.Store)
	ps.Init(dir)

	clog := commit.NewLogger(dir, "mutations", 50<<20)
	clog.Init()
	defer clog.Close()

	ol.init(key, ps, clog)
	glog.Debug("Init successful.")

	edge := x.DirectedEdge{
		Value:     "oh hey there",
		Source:    "new-testing",
		Timestamp: time.Now(),
	}
	if err := ol.AddMutation(edge, Set); err != nil {
		t.Error(err)
	}
	var p types.Posting
	ol.Get(&p, 0)
	if p.Uid() != math.MaxUint64 {
		t.Errorf("All value uids should go to MaxUint64. Got: %v", p.Uid())
	}
	var iout interface{}
	if err := ParseValue(&iout, p.ValueBytes()); err != nil {
		t.Error(err)
	}
	out := iout.(string)
	if out != "oh hey there" {
		t.Errorf("Expected a value. Got: [%q]", out)
	}

	// Run the same check after committing.
	if _, err := ol.MergeIfDirty(); err != nil {
		t.Error(err)
	}
	{
		var tp types.Posting
		if ok := ol.Get(&tp, 0); !ok {
			t.Error("While retrieving posting")
		}
		if err := ParseValue(&iout, tp.ValueBytes()); err != nil {
			t.Error(err)
		}
		out := iout.(string)
		if out != "oh hey there" {
			t.Errorf("Expected a value. Got: [%q]", out)
		}
	}

	// The value made it to the posting list. Changing it now.
	edge.Value = 119
	if err := ol.AddMutation(edge, Set); err != nil {
		t.Error(err)
	}
	if ol.Length() != 1 {
		t.Errorf("Length should be one. Got: %v", ol.Length())
	}
	if ok := ol.Get(&p, 0); !ok {
		t.Error("While retrieving posting")
	}
	if err := ParseValue(&iout, p.ValueBytes()); err != nil {
		t.Error(err)
	}
	intout := iout.(float64)
	if intout != 119 {
		t.Errorf("Expected 119. Got: %v", intout)
	}
}

func TestProcessTask(t *testing.T) {
	// logrus.SetLevel(logrus.DebugLevel)

	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return
	}

	defer os.RemoveAll(dir)
	ps := new(store.Store)
	ps.Init(dir)

	clog := commit.NewLogger(dir, "mutations", 50<<20)
	clog.Init()
	defer clog.Close()

	Init(ps, clog)

	edge := x.DirectedEdge{
		ValueId:   23,
		Source:    "author0",
		Timestamp: time.Now(),
	}
	addEdge(t, edge, GetOrCreate(Key(10, "friend")))
	addEdge(t, edge, GetOrCreate(Key(11, "friend")))
	addEdge(t, edge, GetOrCreate(Key(12, "friend")))

	edge.ValueId = 25
	addEdge(t, edge, GetOrCreate(Key(12, "friend")))

	edge.ValueId = 26
	addEdge(t, edge, GetOrCreate(Key(12, "friend")))

	edge.ValueId = 31
	addEdge(t, edge, GetOrCreate(Key(10, "friend")))
	addEdge(t, edge, GetOrCreate(Key(12, "friend")))

	edge.Value = "photon"
	addEdge(t, edge, GetOrCreate(Key(12, "friend")))

	query := NewQuery("friend", []uint64{10, 11, 12})
	result, err := ProcessTask(query)
	if err != nil {
		t.Error(err)
	}

	ro := flatbuffers.GetUOffsetT(result)
	r := new(task.Result)
	r.Init(result, ro)

	if r.UidmatrixLength() != 3 {
		t.Errorf("Expected 3. Got uidmatrix length: %v", r.UidmatrixLength())
	}
	if err := check(r, 0, []uint64{23, 31}); err != nil {
		t.Error(err)
	}
	if err := check(r, 1, []uint64{23}); err != nil {
		t.Error(err)
	}
	if err := check(r, 2, []uint64{23, 25, 26, 31}); err != nil {
		t.Error(err)
	}

	if r.ValuesLength() != 3 {
		t.Errorf("Expected 3. Got values length: %v", r.ValuesLength())
	}
	var tval task.Value
	if ok := r.Values(&tval, 0); !ok {
		t.Errorf("Unable to retrieve value")
	}
	if tval.ValLength() != 1 ||
		tval.ValBytes()[0] != 0x00 {
		t.Errorf("Invalid byte value at index 0")
	}
	if ok := r.Values(&tval, 1); !ok {
		t.Errorf("Unable to retrieve value")
	}
	if tval.ValLength() != 1 ||
		tval.ValBytes()[0] != 0x00 {
		t.Errorf("Invalid byte value at index 0")
	}

	if ok := r.Values(&tval, 2); !ok {
		t.Errorf("Unable to retrieve value")
	}
	var iout interface{}
	if err := ParseValue(&iout, tval.ValBytes()); err != nil {
		t.Error(err)
	}
	v := iout.(string)
	if v != "photon" {
		t.Errorf("Expected photon. Got: %q", v)
	}
}

func populateGraph(t *testing.T) (string, *store.Store) {
	// logrus.SetLevel(logrus.DebugLevel)
	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return "", nil
	}

	ps := new(store.Store)
	ps.Init(dir)

	worker.Init(ps)

	clog := commit.NewLogger(dir, "mutations", 50<<20)
	clog.Init()
	posting.Init(clog)

	// So, user we're interested in has uid: 1.
	// She has 4 friends: 23, 24, 25, 31, and 101
	edge := x.DirectedEdge{
		ValueId:   23,
		Source:    "testing",
		Timestamp: time.Now(),
	}
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "friend"), ps))

	edge.ValueId = 24
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "friend"), ps))

	edge.ValueId = 25
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "friend"), ps))

	edge.ValueId = 31
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "friend"), ps))

	edge.ValueId = 101
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "friend"), ps))

	// Now let's add a few properties for the main user.
	edge.Value = "Michonne"
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "name"), ps))

	edge.Value = "female"
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "gender"), ps))

	edge.Value = "alive"
	addEdge(t, edge, posting.GetOrCreate(posting.Key(1, "status"), ps))

	// Now let's add a name for each of the friends, except 101.
	edge.Value = "Rick Grimes"
	addEdge(t, edge, posting.GetOrCreate(posting.Key(23, "name"), ps))

	edge.Value = "Glenn Rhee"
	addEdge(t, edge, posting.GetOrCreate(posting.Key(24, "name"), ps))

	edge.Value = "Daryl Dixon"
	addEdge(t, edge, posting.GetOrCreate(posting.Key(25, "name"), ps))

	edge.Value = "Andrea"
	addEdge(t, edge, posting.GetOrCreate(posting.Key(31, "name"), ps))

	return dir, ps
}