remove NewChannelzStorageForTesting entirely

channelz/service/service_sktopt_test.go

@@ -128,8 +128,6 @@ func protoToSocketOption(skopts []*channelzpb.SocketOption) *channelz.SocketOpti
 }
 
 func (s) TestGetSocketOptions(t *testing.T) {
-	channelz.NewChannelzStorageForTesting()
-
 	ss := []*dummySocket{
 		{
 			socketOptions: &channelz.SocketOptionData{

channelz/service/service_test.go

@@ -305,7 +305,6 @@ func (s) TestGetTopChannels(t *testing.T) {
 		},
 		{},
 	}
-	channelz.NewChannelzStorageForTesting()
 
 	for _, c := range tcs {
 		id := channelz.RegisterChannel(c, nil, "")
@@ -358,7 +357,6 @@ func (s) TestGetServers(t *testing.T) {
 			lastCallStartedTimestamp: time.Now().UTC(),
 		},
 	}
-	channelz.NewChannelzStorageForTesting()
 
 	for _, s := range ss {
 		id := channelz.RegisterServer(s, "")
@@ -391,8 +389,6 @@ func (s) TestGetServers(t *testing.T) {
 }
 
 func (s) TestGetServerSockets(t *testing.T) {
-	channelz.NewChannelzStorageForTesting()
-
 	svrID := channelz.RegisterServer(&dummyServer{}, "")
 	defer channelz.RemoveEntry(svrID)
 	refNames := []string{"listen socket 1", "normal socket 1", "normal socket 2"}
@@ -432,8 +428,6 @@ func (s) TestGetServerSockets(t *testing.T) {
 // This test makes a GetServerSockets with a non-zero start ID, and expect only
 // sockets with ID >= the given start ID.
 func (s) TestGetServerSocketsNonZeroStartID(t *testing.T) {
-	channelz.NewChannelzStorageForTesting()
-
 	svrID := channelz.RegisterServer(&dummyServer{}, "")
 	defer channelz.RemoveEntry(svrID)
 	refNames := []string{"listen socket 1", "normal socket 1", "normal socket 2"}
@@ -464,8 +458,6 @@ func (s) TestGetServerSocketsNonZeroStartID(t *testing.T) {
 }
 
 func (s) TestGetChannel(t *testing.T) {
-	channelz.NewChannelzStorageForTesting()
-
 	refNames := []string{"top channel 1", "nested channel 1", "sub channel 2", "nested channel 3"}
 	ids := make([]*channelz.Identifier, 4)
 	ids[0] = channelz.RegisterChannel(&dummyChannel{}, nil, refNames[0])
@@ -577,7 +569,6 @@ func (s) TestGetSubChannel(t *testing.T) {
 		subchanConnectivityChange = fmt.Sprintf("Subchannel Connectivity change to %v", connectivity.Ready)
 		subChanPickNewAddress     = fmt.Sprintf("Subchannel picks a new address %q to connect", "0.0.0.0")
 	)
-	channelz.NewChannelzStorageForTesting()
 
 	refNames := []string{"top channel 1", "sub channel 1", "socket 1", "socket 2"}
 	ids := make([]*channelz.Identifier, 4)
@@ -655,8 +646,6 @@ func (s) TestGetSubChannel(t *testing.T) {
 }
 
 func (s) TestGetSocket(t *testing.T) {
-	channelz.NewChannelzStorageForTesting()
-
 	ss := []*dummySocket{
 		{
 			streamsStarted:                   10,

internal/channelz/funcs.go

@@ -38,8 +38,11 @@ const (
 )
 
 var (
-	db    dbWrapper
-	idGen idGenerator
+	// IDGen is the global channelz entity ID generator.  It should not be used
+	// outside this package except by tests.
+	IDGen IDGenerator
+
+	db dbWrapper
 	// EntryPerPage defines the number of channelz entries to be shown on a web page.
 	EntryPerPage  = int64(50)
 	curState      int32
@@ -50,14 +53,14 @@ var (
 func TurnOn() {
 	if !IsOn() {
 		db.set(newChannelMap())
-		idGen.reset()
+		IDGen.Reset()
 		atomic.StoreInt32(&curState, 1)
 	}
 }
 
 // IsOn returns whether channelz data collection is on.
 func IsOn() bool {
-	return atomic.CompareAndSwapInt32(&curState, 1, 1)
+	return atomic.LoadInt32(&curState) == 1
 }
 
 // SetMaxTraceEntry sets maximum number of trace entry per entity (i.e. channel/subchannel).
@@ -95,13 +98,6 @@ func (d *dbWrapper) get() *channelMap {
 	return d.DB
 }
 
-// NewChannelzStorageForTesting initializes channelz data storage and id
-// generator for testing purposes.
-func NewChannelzStorageForTesting() {
-	db.set(newChannelMap())
-	idGen.reset()
-}
-
 // GetTopChannels returns a slice of top channel's ChannelMetric, along with a
 // boolean indicating whether there's more top channels to be queried for.
 //
@@ -161,7 +157,7 @@ func GetServer(id int64) *ServerMetric {
 //
 // If channelz is not turned ON, the channelz database is not mutated.
 func RegisterChannel(c Channel, pid *Identifier, ref string) *Identifier {
-	id := idGen.genID()
+	id := IDGen.genID()
 	var parent int64
 	isTopChannel := true
 	if pid != nil {
@@ -197,7 +193,7 @@ func RegisterSubChannel(c Channel, pid *Identifier, ref string) (*Identifier, er
 	if pid == nil {
 		return nil, errors.New("a SubChannel's parent id cannot be nil")
 	}
-	id := idGen.genID()
+	id := IDGen.genID()
 	if !IsOn() {
 		return newIdentifer(RefSubChannel, id, pid), nil
 	}
@@ -219,7 +215,7 @@ func RegisterSubChannel(c Channel, pid *Identifier, ref string) (*Identifier, er
 //
 // If channelz is not turned ON, the channelz database is not mutated.
 func RegisterServer(s Server, ref string) *Identifier {
-	id := idGen.genID()
+	id := IDGen.genID()
 	if !IsOn() {
 		return newIdentifer(RefServer, id, nil)
 	}
@@ -245,7 +241,7 @@ func RegisterListenSocket(s Socket, pid *Identifier, ref string) (*Identifier, e
 	if pid == nil {
 		return nil, errors.New("a ListenSocket's parent id cannot be 0")
 	}
-	id := idGen.genID()
+	id := IDGen.genID()
 	if !IsOn() {
 		return newIdentifer(RefListenSocket, id, pid), nil
 	}
@@ -265,7 +261,7 @@ func RegisterNormalSocket(s Socket, pid *Identifier, ref string) (*Identifier, e
 	if pid == nil {
 		return nil, errors.New("a NormalSocket's parent id cannot be 0")
 	}
-	id := idGen.genID()
+	id := IDGen.genID()
 	if !IsOn() {
 		return newIdentifer(RefNormalSocket, id, pid), nil
 	}
@@ -744,14 +740,17 @@ func (c *channelMap) GetServer(id int64) *ServerMetric {
 	return sm
 }
 
-type idGenerator struct {
+// IDGenerator is an incrementing atomic that tracks IDs for channelz entities.
+type IDGenerator struct {
 	id int64
 }
 
-func (i *idGenerator) reset() {
+// Reset resets the generated ID back to zero.  Should only be used at
+// initialization or by tests sensitive to the ID number.
+func (i *IDGenerator) Reset() {
 	atomic.StoreInt64(&i.id, 0)
 }
 
-func (i *idGenerator) genID() int64 {
+func (i *IDGenerator) genID() int64 {
 	return atomic.AddInt64(&i.id, 1)
 }

internal/idle/idle_e2e_test.go

@@ -102,9 +102,6 @@ func channelzTraceEventNotFound(ctx context.Context, wantDesc string) error {
 // Tests the case where channel idleness is disabled by passing an idle_timeout
 // of 0. Verifies that a READY channel with no RPCs does not move to IDLE.
 func (s) TestChannelIdleness_Disabled_NoActivity(t *testing.T) {
-	// Setup channelz for testing.
-	channelz.NewChannelzStorageForTesting()
-
 	// Create a ClientConn with idle_timeout set to 0.
 	r := manual.NewBuilderWithScheme("whatever")
 	dopts := []grpc.DialOption{
@@ -146,9 +143,6 @@ func (s) TestChannelIdleness_Disabled_NoActivity(t *testing.T) {
 // Tests the case where channel idleness is enabled by passing a small value for
 // idle_timeout. Verifies that a READY channel with no RPCs moves to IDLE.
 func (s) TestChannelIdleness_Enabled_NoActivity(t *testing.T) {
-	// Setup channelz for testing.
-	channelz.NewChannelzStorageForTesting()
-
 	// Create a ClientConn with a short idle_timeout.
 	r := manual.NewBuilderWithScheme("whatever")
 	dopts := []grpc.DialOption{
@@ -185,9 +179,6 @@ func (s) TestChannelIdleness_Enabled_NoActivity(t *testing.T) {
 // Tests the case where channel idleness is enabled by passing a small value for
 // idle_timeout. Verifies that a READY channel with an ongoing RPC stays READY.
 func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
-	// Setup channelz for testing.
-	channelz.NewChannelzStorageForTesting()
-
 	// Create a ClientConn with a short idle_timeout.
 	r := manual.NewBuilderWithScheme("whatever")
 	dopts := []grpc.DialOption{
@@ -269,9 +260,6 @@ func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
 // idle_timeout. Verifies that activity on a READY channel (frequent and short
 // RPCs) keeps it from moving to IDLE.
 func (s) TestChannelIdleness_Enabled_ActiveSinceLastCheck(t *testing.T) {
-	// Setup channelz for testing.
-	channelz.NewChannelzStorageForTesting()
-
 	// Create a ClientConn with a short idle_timeout.
 	r := manual.NewBuilderWithScheme("whatever")
 	dopts := []grpc.DialOption{
@@ -332,9 +320,6 @@ func (s) TestChannelIdleness_Enabled_ActiveSinceLastCheck(t *testing.T) {
 // idle_timeout. Verifies that a READY channel with no RPCs moves to IDLE. Also
 // verifies that a subsequent RPC on the IDLE channel kicks it out of IDLE.
 func (s) TestChannelIdleness_Enabled_ExitIdleOnRPC(t *testing.T) {
-	// Setup channelz for testing.
-	channelz.NewChannelzStorageForTesting()
-
 	// Start a test backend and set the bootstrap state of the resolver to
 	// include this address. This will ensure that when the resolver is
 	// restarted when exiting idle, it will push the same address to grpc again.
@@ -393,9 +378,6 @@ func (s) TestChannelIdleness_Enabled_ExitIdleOnRPC(t *testing.T) {
 //
 // In either of these cases, all RPCs must succeed.
 func (s) TestChannelIdleness_Enabled_IdleTimeoutRacesWithRPCs(t *testing.T) {
-	// Setup channelz for testing.
-	channelz.NewChannelzStorageForTesting()
-
 	// Start a test backend and set the bootstrap state of the resolver to
 	// include this address. This will ensure that when the resolver is
 	// restarted when exiting idle, it will push the same address to grpc again.

internal/testutils/wrappers.go

@@ -72,13 +72,3 @@ func NewListenerWrapper(t *testing.T, lis net.Listener) *ListenerWrapper {
 		NewConnCh: NewChannel(),
 	}
 }
-
-// ErrCloseWrapper wraps closer with a function that does not return an error,
-// but calls t.Error if it does, instead.
-func ErrCloseWrapper(t *testing.T, closer func() error) func() {
-	return func() {
-		if err := closer(); err != nil {
-			t.Error(err)
-		}
-	}
-}

test/balancer_switching_test.go

@@ -29,7 +29,6 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/balancer/stub"
-	"google.golang.org/grpc/internal/channelz"
 	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/internal/testutils/fakegrpclb"
 	"google.golang.org/grpc/internal/testutils/pickfirst"
@@ -63,7 +62,6 @@ const (
 //
 // Returns a cleanup function to be invoked by the caller.
 func setupBackendsAndFakeGRPCLB(t *testing.T) ([]*stubserver.StubServer, *fakegrpclb.Server, func()) {
-	channelz.NewChannelzStorageForTesting()
 	backends, backendsCleanup := startBackendsForBalancerSwitch(t)
 
 	lbServer, err := fakegrpclb.NewServer(fakegrpclb.ServerParams{

test/channelz_linux_test.go

@@ -35,8 +35,6 @@ func (s) TestCZSocketMetricsSocketOption(t *testing.T) {
 }
 
 func testCZSocketMetricsSocketOption(t *testing.T, e env) {
-	channelz.NewChannelzStorageForTesting()
-
 	te := newTest(t, e)
 	te.startServer(&testServer{security: e.security})
 	defer te.tearDown()