Golang topo.Impl类代码示例

// waitForInitialValue waits for the initial value of
// /keyspaces/test_keyspace/SrvKeyspace to appear, and match the
// provided srvKeyspace.
func waitForInitialValue(t *testing.T, ts topo.Impl, cell string, srvKeyspace *topodatapb.SrvKeyspace) (changes <-chan *topo.WatchData, cancel func()) {
	var current *topo.WatchData
	ctx := context.Background()
	start := time.Now()
	for {
		current, changes, cancel = ts.Watch(ctx, cell, "/keyspaces/test_keyspace/SrvKeyspace")
		if current.Err == topo.ErrNoNode {
			// hasn't appeared yet
			if time.Now().Sub(start) > 10*time.Second {
				t.Fatalf("time out waiting for file to appear")
			}
			time.Sleep(10 * time.Millisecond)
			continue
		}
		if current.Err != nil {
			t.Fatalf("watch failed: %v", current.Err)
		}
		// we got a valid result
		break
	}
	got := &topodatapb.SrvKeyspace{}
	if err := proto.Unmarshal(current.Contents, got); err != nil {
		t.Fatalf("cannot proto-unmarshal data: %v", err)
	}
	if !proto.Equal(got, srvKeyspace) {
		t.Fatalf("got bad data: %v expected: %v", got, srvKeyspace)
	}

	return changes, cancel
}

func checkKeyspaceLockTimeout(ctx context.Context, t *testing.T, ts topo.Impl) {
	lockPath, err := ts.LockKeyspaceForAction(ctx, "test_keyspace", "fake-content")
	if err != nil {
		t.Fatalf("LockKeyspaceForAction: %v", err)
	}

	// test we can't take the lock again
	fastCtx, cancel := context.WithTimeout(ctx, timeUntilLockIsTaken)
	if _, err := ts.LockKeyspaceForAction(fastCtx, "test_keyspace", "unused-fake-content"); err != topo.ErrTimeout {
		t.Fatalf("LockKeyspaceForAction(again): %v", err)
	}
	cancel()

	// test we can interrupt taking the lock
	interruptCtx, cancel := context.WithCancel(ctx)
	go func() {
		time.Sleep(timeUntilLockIsTaken)
		cancel()
	}()
	if _, err := ts.LockKeyspaceForAction(interruptCtx, "test_keyspace", "unused-fake-content"); err != topo.ErrInterrupted {
		t.Fatalf("LockKeyspaceForAction(interrupted): %v", err)
	}

	if err := ts.UnlockKeyspaceForAction(ctx, "test_keyspace", lockPath, "fake-results"); err != nil {
		t.Fatalf("UnlockKeyspaceForAction(): %v", err)
	}

	// test we can't unlock again
	if err := ts.UnlockKeyspaceForAction(ctx, "test_keyspace", lockPath, "fake-results"); err == nil {
		t.Fatalf("UnlockKeyspaceForAction(again) worked")
	}
}

// CheckKeyspaceLock checks we can take a keyspace lock as expected.
func CheckKeyspaceLock(ctx context.Context, t *testing.T, ts topo.Impl) {
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &pb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}

	checkKeyspaceLockTimeout(ctx, t, ts)
	checkKeyspaceLockMissing(ctx, t, ts)
	checkKeyspaceLockUnblocks(ctx, t, ts)
}

func getLocalCell(ctx context.Context, t *testing.T, ts topo.Impl) string {
	cells, err := ts.GetKnownCells(ctx)
	if err != nil {
		t.Fatalf("GetKnownCells: %v", err)
	}
	if len(cells) < 1 {
		t.Fatalf("provided topo.Impl doesn't have enough cells (need at least 1): %v", cells)
	}
	return cells[0]
}

// CopyShardReplications will create the ShardReplication objects in
// the destination topo
func CopyShardReplications(ctx context.Context, fromTS, toTS topo.Impl) {
	keyspaces, err := fromTS.GetKeyspaces(ctx)
	if err != nil {
		log.Fatalf("fromTS.GetKeyspaces: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()
			shards, err := fromTS.GetShardNames(ctx, keyspace)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetShardNames(%v): %v", keyspace, err))
				return
			}

			for _, shard := range shards {
				wg.Add(1)
				go func(keyspace, shard string) {
					defer wg.Done()

					// read the source shard to get the cells
					s, _, err := fromTS.GetShard(ctx, keyspace, shard)
					if err != nil {
						rec.RecordError(fmt.Errorf("GetShard(%v, %v): %v", keyspace, shard, err))
						return
					}

					for _, cell := range s.Cells {
						sri, err := fromTS.GetShardReplication(ctx, cell, keyspace, shard)
						if err != nil {
							rec.RecordError(fmt.Errorf("GetShardReplication(%v, %v, %v): %v", cell, keyspace, shard, err))
							continue
						}

						if err := toTS.UpdateShardReplicationFields(ctx, cell, keyspace, shard, func(oldSR *topodatapb.ShardReplication) error {
							*oldSR = *sri.ShardReplication
							return nil
						}); err != nil {
							rec.RecordError(fmt.Errorf("UpdateShardReplicationFields(%v, %v, %v): %v", cell, keyspace, shard, err))
						}
					}
				}(keyspace, shard)
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyShards failed: %v", rec.Error())
	}
}

// CopyTablets will create the tablets in the destination topo
func CopyTablets(ctx context.Context, fromTS, toTS topo.Impl) {
	cells, err := fromTS.GetKnownCells(ctx)
	if err != nil {
		log.Fatalf("fromTS.GetKnownCells: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, cell := range cells {
		wg.Add(1)
		go func(cell string) {
			defer wg.Done()
			tabletAliases, err := fromTS.GetTabletsByCell(ctx, cell)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetTabletsByCell(%v): %v", cell, err))
			} else {
				for _, tabletAlias := range tabletAliases {
					wg.Add(1)
					go func(tabletAlias *topodatapb.TabletAlias) {
						defer wg.Done()

						// read the source tablet
						tablet, _, err := fromTS.GetTablet(ctx, tabletAlias)
						if err != nil {
							rec.RecordError(fmt.Errorf("GetTablet(%v): %v", tabletAlias, err))
							return
						}

						// try to create the destination
						err = toTS.CreateTablet(ctx, tablet)
						if err == topo.ErrNodeExists {
							// update the destination tablet
							log.Warningf("tablet %v already exists, updating it", tabletAlias)
							_, err = toTS.UpdateTabletFields(ctx, tablet.Alias, func(t *topodatapb.Tablet) error {
								*t = *tablet
								return nil
							})
						}
						if err != nil {
							rec.RecordError(fmt.Errorf("CreateTablet(%v): %v", tabletAlias, err))
							return
						}
					}(tabletAlias)
				}
			}
		}(cell)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyTablets failed: %v", rec.Error())
	}
}

// checkWatchInterrupt tests we can interrupt a watch.
func checkWatchInterrupt(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	cell := getLocalCell(ctx, t, ts)

	// create some data
	srvKeyspace := &topodatapb.SrvKeyspace{
		ShardingColumnName: "user_id",
	}
	if err := ts.UpdateSrvKeyspace(ctx, cell, "test_keyspace", srvKeyspace); err != nil {
		t.Fatalf("UpdateSrvKeyspace(1): %v", err)
	}

	// Start watching, it should work.
	changes, cancel := waitForInitialValue(t, ts, cell, srvKeyspace)

	// Now cancel the watch.
	cancel()

	// Make sure we get the topo.ErrInterrupted notification eventually.
	for {
		wd, ok := <-changes
		if !ok {
			t.Fatalf("watch channel unexpectedly closed")
		}
		if wd.Err == topo.ErrInterrupted {
			// good
			break
		}
		if wd.Err != nil {
			t.Fatalf("bad error returned for deletion: %v", wd.Err)
		}
		// we got something, better be the right value
		got := &topodatapb.SrvKeyspace{}
		if err := proto.Unmarshal(wd.Contents, got); err != nil {
			t.Fatalf("cannot proto-unmarshal data: %v", err)
		}
		if got.ShardingColumnName == "user_id" {
			// good value
			continue
		}
		t.Fatalf("got unknown SrvKeyspace waiting for deletion: %v", got)
	}

	// Now the channel should be closed.
	if wd, ok := <-changes; ok {
		t.Fatalf("got unexpected event after error: %v", wd)
	}

	// And calling cancel() again should just work.
	cancel()
}

// CheckShardLock checks we can take a shard lock
func CheckShardLock(ctx context.Context, t *testing.T, ts topo.Impl) {
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &pb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "10-20", &pb.Shard{
		KeyRange: newKeyRange3("10-20"),
	}); err != nil {
		t.Fatalf("CreateShard: %v", err)
	}

	checkShardLockTimeout(ctx, t, ts)
	checkShardLockMissing(ctx, t, ts)
	checkShardLockUnblocks(ctx, t, ts)
}

func checkListDir(ctx context.Context, t *testing.T, ts topo.Impl, cell string, dirPath string, expected []string) {
	entries, err := ts.ListDir(ctx, cell, dirPath)
	switch err {
	case topo.ErrNoNode:
		if len(expected) != 0 {
			t.Errorf("ListDir(%v) returned ErrNoNode but was expecting %v", dirPath, expected)
		}
	case nil:
		if !reflect.DeepEqual(entries, expected) {
			t.Errorf("ListDir(%v) returned %v but was expecting %v", dirPath, entries, expected)
		}
	default:
		t.Errorf("ListDir(%v) returned unexpected error: %v", dirPath, err)
	}
}

// checkKeyspaceLock checks we can take a keyspace lock as expected.
func checkKeyspaceLock(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}

	t.Log("===      checkKeyspaceLockTimeout")
	checkKeyspaceLockTimeout(ctx, t, ts)

	t.Log("===      checkKeyspaceLockMissing")
	checkKeyspaceLockMissing(ctx, t, ts)

	t.Log("===      checkKeyspaceLockUnblocks")
	checkKeyspaceLockUnblocks(ctx, t, ts)
}

// checkShardLock checks we can take a shard lock
func checkShardLock(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	if err := ts.CreateKeyspace(ctx, "test_keyspace", &topodatapb.Keyspace{}); err != nil {
		t.Fatalf("CreateKeyspace: %v", err)
	}
	if err := ts.CreateShard(ctx, "test_keyspace", "10-20", &topodatapb.Shard{
		KeyRange: newKeyRange("10-20"),
	}); err != nil {
		t.Fatalf("CreateShard: %v", err)
	}

	t.Log("===     checkShardLockTimeout")
	checkShardLockTimeout(ctx, t, ts)

	t.Log("===     checkShardLockMissing")
	checkShardLockMissing(ctx, t, ts)

	t.Log("===     checkShardLockUnblocks")
	checkShardLockUnblocks(ctx, t, ts)
}

// CopyKeyspaces will create the keyspaces in the destination topo
func CopyKeyspaces(ctx context.Context, fromTS, toTS topo.Impl) {
	keyspaces, err := fromTS.GetKeyspaces(ctx)
	if err != nil {
		log.Fatalf("GetKeyspaces: %v", err)
	}

	wg := sync.WaitGroup{}
	rec := concurrency.AllErrorRecorder{}
	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			defer wg.Done()

			k, _, err := fromTS.GetKeyspace(ctx, keyspace)
			if err != nil {
				rec.RecordError(fmt.Errorf("GetKeyspace(%v): %v", keyspace, err))
				return
			}

			if err := toTS.CreateKeyspace(ctx, keyspace, k); err != nil {
				if err == topo.ErrNodeExists {
					log.Warningf("keyspace %v already exists", keyspace)
				} else {
					rec.RecordError(fmt.Errorf("CreateKeyspace(%v): %v", keyspace, err))
				}
			}
		}(keyspace)
	}
	wg.Wait()
	if rec.HasErrors() {
		log.Fatalf("copyKeyspaces failed: %v", rec.Error())
	}
}

// checkSrvShardLockUnblocks makes sure that a routine waiting on a lock
// is unblocked when another routine frees the lock
func checkSrvShardLockUnblocks(ctx context.Context, t *testing.T, ts topo.Impl) {
	cell := getLocalCell(ctx, t, ts)
	unblock := make(chan struct{})
	finished := make(chan struct{})

	// as soon as we're unblocked, we try to lock the shard
	go func() {
		<-unblock
		lockPath, err := ts.LockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", "fake-content")
		if err != nil {
			t.Fatalf("LockSrvShardForAction(test, test_keyspace, 10-20) failed: %v", err)
		}
		if err = ts.UnlockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", lockPath, "fake-results"); err != nil {
			t.Fatalf("UnlockSrvShardForAction(test, test_keyspace, 10-20): %v", err)
		}
		close(finished)
	}()

	// lock the shard
	lockPath2, err := ts.LockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", "fake-content")
	if err != nil {
		t.Fatalf("LockSrvShardForAction(test, test_keyspace, 10-20) failed: %v", err)
	}

	// unblock the go routine so it starts waiting
	close(unblock)

	// sleep for a while so we're sure the go routine is blocking
	time.Sleep(timeUntilLockIsTaken)

	if err = ts.UnlockSrvShardForAction(ctx, cell, "test_keyspace", "10-20", lockPath2, "fake-results"); err != nil {
		t.Fatalf("UnlockSrvShardForAction(test, test_keyspace, 10-20): %v", err)
	}

	timeout := time.After(10 * time.Second)
	select {
	case <-finished:
	case <-timeout:
		t.Fatalf("unlocking timed out")
	}
}

// checkSrvVSchema tests the SrvVSchema methods (other than watch).
func checkSrvVSchema(t *testing.T, ts topo.Impl) {
	ctx := context.Background()
	cell := getLocalCell(ctx, t, ts)

	// check GetSrvVSchema returns topo.ErrNoNode if no SrvVSchema
	if _, err := ts.GetSrvVSchema(ctx, cell); err != topo.ErrNoNode {
		t.Errorf("GetSrvVSchema(not set): %v", err)
	}

	srvVSchema := &vschemapb.SrvVSchema{
		Keyspaces: map[string]*vschemapb.Keyspace{
			"test_keyspace": {
				Sharded: true,
			},
		},
	}
	if err := ts.UpdateSrvVSchema(ctx, cell, srvVSchema); err != nil {
		t.Errorf("UpdateSrvVSchema(1): %v", err)
	}
	if v, err := ts.GetSrvVSchema(ctx, cell); err != nil || !proto.Equal(srvVSchema, v) {
		t.Errorf("GetSrvVSchema(valid): %v %v", err, v)
	}
}

func checkShardLockMissing(ctx context.Context, t *testing.T, ts topo.Impl) {
	// test we can't lock a non-existing shard
	if _, err := ts.LockShardForAction(ctx, "test_keyspace", "20-30", "fake-content"); err == nil {
		t.Fatalf("LockShardForAction(test_keyspace/20-30) worked for non-existing shard")
	}
}