Golang x.DataKey函数代码示例

func TestGenerateGroup(t *testing.T) {
	dir, err := ioutil.TempDir("", "store3")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	r := bytes.NewReader([]byte("default: fp % 3"))
	require.NoError(t, group.ParseConfig(r), "Unable to parse config.")

	ps, err := store.NewStore(dir)
	if err != nil {
		t.Fatal(err)
	}
	defer ps.Close()
	posting.Init(ps)
	Init(ps)

	require.Equal(t, uint32(0), group.BelongsTo("pred0"))
	writePLs(t, "pred0", 33, 1, ps)

	require.Equal(t, uint32(1), group.BelongsTo("p1"))
	writePLs(t, "p1", 34, 1, ps)

	require.Equal(t, uint32(2), group.BelongsTo("pr2"))
	writePLs(t, "pr2", 35, 1, ps)
	time.Sleep(time.Second)

	g, err := generateGroup(0)
	if err != nil {
		t.Error(err)
	}
	require.Equal(t, 33, len(g.Keys))
	for i, k := range g.Keys {
		require.Equal(t, x.DataKey("pred0", uint64(i)), k.Key)
	}

	g, err = generateGroup(1)
	if err != nil {
		t.Error(err)
	}
	require.Equal(t, 34, len(g.Keys))
	for i, k := range g.Keys {
		require.Equal(t, x.DataKey("p1", uint64(i)), k.Key)
	}

	g, err = generateGroup(2)
	if err != nil {
		t.Error(err)
	}
	require.Equal(t, 35, len(g.Keys))
	for i, k := range g.Keys {
		require.Equal(t, x.DataKey("pr2", uint64(i)), k.Key)
	}
}

func BenchmarkAddMutations(b *testing.B) {
	// logrus.SetLevel(logrus.DebugLevel)
	key := x.DataKey("name", 1)
	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		b.Error(err)
		return
	}

	defer os.RemoveAll(dir)
	ps, err := store.NewStore(dir)
	if err != nil {
		b.Error(err)
		return
	}

	l := getNew(key, ps)
	b.ResetTimer()

	ctx := context.Background()
	for i := 0; i < b.N; i++ {
		if err != nil {
			b.Error(err)
			return
		}
		edge := &task.DirectedEdge{
			ValueId: uint64(rand.Intn(b.N) + 1),
			Label:   "testing",
		}
		if _, err := l.AddMutation(ctx, edge, Set); err != nil {
			b.Error(err)
		}
	}
}

func TestAddMutation_Value(t *testing.T) {
	key := x.DataKey("value", 10)
	dir, err := ioutil.TempDir("", "storetest_")
	if err != nil {
		t.Error(err)
		return
	}
	defer os.RemoveAll(dir)

	ps, err := store.NewStore(dir)
	require.NoError(t, err)
	Init(ps)

	ol := getNew(key, ps)
	log.Println("Init successful.")

	edge := &task.DirectedEdge{
		Value: []byte("oh hey there"),
		Label: "new-testing",
	}
	addMutation(t, ol, edge, Set)
	checkValue(t, ol, "oh hey there")

	// Run the same check after committing.
	_, err = ol.CommitIfDirty(context.Background())
	require.NoError(t, err)
	checkValue(t, ol, "oh hey there")

	// The value made it to the posting list. Changing it now.
	edge.Value = []byte(strconv.Itoa(119))
	addMutation(t, ol, edge, Set)
	checkValue(t, ol, "119")
}

func TestAfterUIDCount2(t *testing.T) {
	key := x.DataKey("value", 10)
	dir, err := ioutil.TempDir("", "storetest_")
	require.NoError(t, err)
	defer os.RemoveAll(dir)

	ps, err := store.NewStore(dir)
	require.NoError(t, err)
	ol := getNew(key, ps)

	// Set value to cars and merge to RocksDB.
	edge := &task.DirectedEdge{
		Label: "jchiu",
	}

	for i := 100; i < 300; i++ {
		edge.ValueId = uint64(i)
		addMutation(t, ol, edge, Set)
	}
	require.EqualValues(t, 200, ol.Length(0))
	require.EqualValues(t, 100, ol.Length(199))
	require.EqualValues(t, 0, ol.Length(300))

	// Re-insert 1/4 of the edges. Counts should not change.
	edge.Label = "somethingelse"
	for i := 100; i < 300; i += 4 {
		edge.ValueId = uint64(i)
		addMutation(t, ol, edge, Set)
	}
	require.EqualValues(t, 200, ol.Length(0))
	require.EqualValues(t, 100, ol.Length(199))
	require.EqualValues(t, 0, ol.Length(300))
}

func populateGraphForSort(t *testing.T, ps *store.Store) {
	edge := &task.DirectedEdge{
		Label: "author1",
		Attr:  "dob",
	}

	dobs := []string{
		"1980-05-05", // 10 (1980)
		"1980-04-05", // 11
		"1979-05-05", // 12 (1979)
		"1979-02-05", // 13
		"1979-03-05", // 14
		"1965-05-05", // 15 (1965)
		"1965-04-05", // 16
		"1965-03-05", // 17
		"1970-05-05", // 18 (1970)
		"1970-04-05", // 19
		"1970-01-05", // 20
		"1970-02-05", // 21
	}
	// The sorted UIDs are: (17 16 15) (20 21 19 18) (13 14 12) (11 10)

	for i, dob := range dobs {
		edge.Entity = uint64(i + 10)
		edge.Value = []byte(dob)
		addEdge(t, edge,
			getOrCreate(x.DataKey(edge.Attr, edge.Entity)))
	}
	time.Sleep(200 * time.Millisecond) // Let indexing finish.
}

// allocateUniqueUid returns an integer in range:
// [minIdx, maxIdx] derived based on numInstances and instanceIdx.
// which hasn't already been allocated to other xids. It does this by
// taking the fingerprint of the xid appended with zero or more spaces
// until the obtained integer is unique.
func allocateUniqueUid(instanceIdx uint64, numInstances uint64) uint64 {
	mod := math.MaxUint64 / numInstances
	minIdx := instanceIdx * mod

	buf := make([]byte, 128)
	for {
		_, err := rand.Read(buf)
		x.Checkf(err, "rand.Read shouldn't throw an error")

		uidb := farm.Fingerprint64(buf) // Generate from hash.
		uid := (uidb % mod) + minIdx
		if uid == math.MaxUint64 || !lmgr.isNew(uid) {
			continue
		}

		// Check if this uid has already been allocated.
		key := x.DataKey("_uid_", uid)
		pl, decr := posting.GetOrCreate(key)
		defer decr()

		if pl.Length(0) == 0 {
			return uid
		}
	}
	log.Fatalf("This shouldn't be reached.")
	return 0
}

func TestAddMutation_jchiu3(t *testing.T) {
	key := x.DataKey("value", 10)
	dir, err := ioutil.TempDir("", "storetest_")
	require.NoError(t, err)
	defer os.RemoveAll(dir)

	ps, err := store.NewStore(dir)
	require.NoError(t, err)
	Init(ps)
	ol := getNew(key, ps)

	// Set value to cars and merge to RocksDB.
	edge := &task.DirectedEdge{
		Value: []byte("cars"),
		Label: "jchiu",
	}
	addMutation(t, ol, edge, Set)
	require.Equal(t, 1, ol.Length(0))
	merged, err := ol.CommitIfDirty(context.Background())
	require.NoError(t, err)
	require.True(t, merged)
	require.EqualValues(t, 1, ol.Length(0))
	checkValue(t, ol, "cars")

	// Del a value cars and but don't merge.
	edge = &task.DirectedEdge{
		Value: []byte("cars"),
		Label: "jchiu",
	}
	addMutation(t, ol, edge, Del)
	require.Equal(t, 0, ol.Length(0))

	// Set value to newcars, but don't merge yet.
	edge = &task.DirectedEdge{
		Value: []byte("newcars"),
		Label: "jchiu",
	}
	addMutation(t, ol, edge, Set)
	require.EqualValues(t, 1, ol.Length(0))
	checkValue(t, ol, "newcars")

	// Del a value othercars and but don't merge.
	edge = &task.DirectedEdge{
		Value: []byte("othercars"),
		Label: "jchiu",
	}
	addMutation(t, ol, edge, Del)
	require.NotEqual(t, 0, ol.Length(0))
	checkValue(t, ol, "newcars")

	// Del a value newcars and but don't merge.
	edge = &task.DirectedEdge{
		Value: []byte("newcars"),
		Label: "jchiu",
	}
	addMutation(t, ol, edge, Del)
	require.Equal(t, 0, ol.Length(0))
}

func TestAddMutation_gru2(t *testing.T) {
	key := x.DataKey("question.tag", 0x01)
	dir, err := ioutil.TempDir("", "storetest_")
	require.NoError(t, err)
	defer os.RemoveAll(dir)

	ps, err := store.NewStore(dir)
	require.NoError(t, err)
	Init(ps)
	ol := getNew(key, ps)

	{
		// Set two tag ids and merge.
		edge := &task.DirectedEdge{
			ValueId: 0x02,
			Label:   "gru",
		}
		addMutation(t, ol, edge, Set)
		edge = &task.DirectedEdge{
			ValueId: 0x03,
			Label:   "gru",
		}
		addMutation(t, ol, edge, Set)
		merged, err := ol.CommitIfDirty(context.Background())
		require.NoError(t, err)
		require.True(t, merged)
	}

	{
		// Lets set a new tag and delete the two older ones.
		edge := &task.DirectedEdge{
			ValueId: 0x02,
			Label:   "gru",
		}
		addMutation(t, ol, edge, Del)
		edge = &task.DirectedEdge{
			ValueId: 0x03,
			Label:   "gru",
		}
		addMutation(t, ol, edge, Del)

		edge = &task.DirectedEdge{
			ValueId: 0x04,
			Label:   "gru",
		}
		addMutation(t, ol, edge, Set)

		merged, err := ol.CommitIfDirty(context.Background())
		require.NoError(t, err)
		require.True(t, merged)
	}

	// Posting list should just have the new tag.
	uids := []uint64{0x04}
	require.Equal(t, uids, listToArray(t, 0, ol))
}

func addEdgeToUID(t *testing.T, ps *store.Store, attr string, src uint64, dst uint64) {
	edge := &task.DirectedEdge{
		ValueId: dst,
		Label:   "testing",
		Attr:    attr,
		Entity:  src,
	}
	l, _ := posting.GetOrCreate(x.DataKey(attr, src))
	require.NoError(t,
		l.AddMutationWithIndex(context.Background(), edge, posting.Set))
}

func markTaken(ctx context.Context, uid uint64) {
	mu := &task.DirectedEdge{
		Entity: uid,
		Attr:   "_uid_",
		Value:  []byte("_"), // not txid
		Label:  "_loader_",
	}
	key := x.DataKey("_uid_", uid)
	plist, decr := posting.GetOrCreate(key)
	plist.AddMutation(ctx, mu, posting.Set)
	decr()
}

func populateGraphBackup(t *testing.T) {
	edge := &task.DirectedEdge{
		ValueId: 5,
		Label:   "author0",
		Attr:    "friend",
	}
	edge.Entity = 1
	addEdge(t, edge, getOrCreate(x.DataKey("friend", 1)))

	edge.Entity = 2
	addEdge(t, edge, getOrCreate(x.DataKey("friend", 2)))

	edge.Entity = 3
	addEdge(t, edge, getOrCreate(x.DataKey("friend", 3)))

	edge.Entity = 4
	addEdge(t, edge, getOrCreate(x.DataKey("friend", 4)))

	edge.Entity = 1
	edge.ValueId = 0
	edge.Value = []byte("pho\\ton")
	edge.Attr = "name"
	addEdge(t, edge, getOrCreate(x.DataKey("name", 1)))

	edge.Entity = 2
	addEdge(t, edge, getOrCreate(x.DataKey("name", 2)))
}

func addEdgeToTypedValue(t *testing.T, ps *store.Store, attr string, src uint64,
	typ types.TypeID, value []byte) {
	edge := &task.DirectedEdge{
		Value:     value,
		ValueType: uint32(typ),
		Label:     "testing",
		Attr:      attr,
		Entity:    src,
	}
	l, _ := posting.GetOrCreate(x.DataKey(attr, src))
	require.NoError(t,
		l.AddMutationWithIndex(context.Background(), edge, posting.Set))
}

// handleNQuads converts the nQuads that satisfy the modulo
// rule into posting lists.
func (s *state) handleNQuads(wg *sync.WaitGroup) {
	defer wg.Done()
	// Check if we need to mark used UIDs.
	markUids := s.groupsMap[group.BelongsTo("_uid_")]
	ctx := context.Background()
	for nq := range s.cnq {
		if s.Error() != nil {
			return
		}
		// Only handle this edge if the attribute satisfies the modulo rule
		if !s.groupsMap[group.BelongsTo(nq.Predicate)] {
			atomic.AddUint64(&s.ctr.ignored, 1)
			continue
		}

		edge, err := nq.ToEdge()
		for err != nil {
			// Just put in a retry loop to tackle temporary errors.
			if err == posting.ErrRetry {
				time.Sleep(time.Microsecond)

			} else {
				s.SetError(err)
				glog.WithError(err).WithField("nq", nq).
					Error("While converting to edge")
				return
			}
			edge, err = nq.ToEdge()
		}

		key := x.DataKey(edge.Attr, edge.Entity)

		plist, decr := posting.GetOrCreate(key)
		plist.AddMutationWithIndex(ctx, edge, posting.Set)
		decr() // Don't defer, just call because we're in a channel loop.

		// Mark UIDs and XIDs as taken
		if markUids {
			// Mark entity UID.
			markTaken(ctx, edge.Entity)
			// Mark the Value UID.
			if edge.ValueId != 0 {
				markTaken(ctx, edge.ValueId)
			}
		}
		atomic.AddUint64(&s.ctr.processed, 1)
	}
}

func writePLs(t *testing.T, pred string, count int, vid uint64, ps *store.Store) {
	for i := 0; i < count; i++ {
		k := x.DataKey(pred, uint64(i))
		list, _ := posting.GetOrCreate(k)

		de := &task.DirectedEdge{
			ValueId: vid,
			Label:   "test",
		}
		list.AddMutation(context.TODO(), de, posting.Set)
		if merged, err := list.CommitIfDirty(context.TODO()); err != nil {
			t.Errorf("While merging: %v", err)
		} else if !merged {
			t.Errorf("No merge happened")
		}
	}
}

// runMutations goes through all the edges and applies them. It returns the
// mutations which were not applied in left.
func runMutations(ctx context.Context, edges []*task.DirectedEdge, op uint32) error {
	for _, edge := range edges {
		if !groups().ServesGroup(group.BelongsTo(edge.Attr)) {
			return x.Errorf("Predicate fingerprint doesn't match this instance")
		}

		key := x.DataKey(edge.Attr, edge.Entity)
		plist, decr := posting.GetOrCreate(key)
		defer decr()

		if err := plist.AddMutationWithIndex(ctx, edge, op); err != nil {
			x.Printf("Error while adding mutation: %v %v", edge, err)
			return err // abort applying the rest of them.
		}
	}
	return nil
}