diff --git a/weed/server/master_grpc_server.go b/weed/server/master_grpc_server.go index d45ee860e..15cac59f6 100644 --- a/weed/server/master_grpc_server.go +++ b/weed/server/master_grpc_server.go @@ -418,7 +418,7 @@ func (ms *MasterServer) KeepConnected(stream master_pb.Seaweed_KeepConnectedServ if req.ClientType == cluster.FilerType { ms.LockRingManager.RemoveServer(cluster.FilerGroupName(req.FilerGroup), peerAddress) } - ms.deleteClient(clientName) + ms.deleteClient(clientName, messageChan) }() // Send volume locations to the client @@ -481,7 +481,12 @@ func (ms *MasterServer) KeepConnected(stream master_pb.Seaweed_KeepConnectedServ defer ticker.Stop() for { select { - case message := <-messageChan: + case message, ok := <-messageChan: + if !ok { + // a closed channel receives nil forever; without this check the + // loop would flood the client with empty messages at wire speed + return nil + } if err := stream.Send(message); err != nil { // The error, not the message: it carries every volume id on a // newly connected node, and formatting a proto that size to @@ -582,14 +587,18 @@ func (ms *MasterServer) addClient(filerGroup, clientType string, clientAddress p return } -func (ms *MasterServer) deleteClient(clientName string) { +func (ms *MasterServer) deleteClient(clientName string, messageChan chan *master_pb.KeepConnectedResponse) { glog.V(0).Infof("- client %v", clientName) ms.clientChansLock.Lock() - // close message chan, so that the KeepConnected go routine can exit - if clientChan, ok := ms.clientChans[clientName]; ok { - close(clientChan) + // a client that reconnects before the old handler exits re-registers the + // same name, so the map may already hold the new stream's channel: close + // only our own, and leave the entry alone unless it is still ours + if ms.clientChans[clientName] == messageChan { delete(ms.clientChans, clientName) } + // safe under the write lock: broadcasters send while holding the read + // lock and only to channels found in the map + close(messageChan) ms.clientChansLock.Unlock() } diff --git a/weed/server/master_grpc_server_test.go b/weed/server/master_grpc_server_test.go index fb12f5dc9..7822078a4 100644 --- a/weed/server/master_grpc_server_test.go +++ b/weed/server/master_grpc_server_test.go @@ -38,6 +38,50 @@ func TestInitialLockRingUpdateSkipsNonFilers(t *testing.T) { assert.Nil(t, ms.initialLockRingUpdate(cluster.BrokerType, "group-a")) } +// TestReconnectedClientSurvivesOldHandlerCleanup covers a client reconnecting +// KeepConnected under the same name before the old handler has exited: the old +// handler's deferred cleanup must not close the channel the reconnected stream +// registered, or that stream reads nil from its closed channel forever and +// floods the client with empty responses. +func TestReconnectedClientSurvivesOldHandlerCleanup(t *testing.T) { + ms := &MasterServer{ + clientChans: make(map[string]chan *master_pb.KeepConnectedResponse), + } + + clientName, oldChan := ms.addClient("", cluster.MasterType, "peer:19333") + _, newChan := ms.addClient("", cluster.MasterType, "peer:19333") + + // the old handler's deferred cleanup runs after the reconnect registered + ms.deleteClient(clientName, oldChan) + + // the old channel is closed so its drain goroutine can exit + select { + case _, ok := <-oldChan: + assert.False(t, ok, "old channel should be closed") + default: + t.Fatal("old channel left open") + } + + // the reconnected stream's channel is untouched and still receives broadcasts + ms.broadcastToClients(&master_pb.KeepConnectedResponse{ + VolumeLocation: &master_pb.VolumeLocation{Url: "volume-a:8080"}, + }) + select { + case message, ok := <-newChan: + require.True(t, ok, "reconnected client's channel was closed by the old handler's cleanup") + require.NotNil(t, message.GetVolumeLocation()) + default: + t.Fatal("no broadcast reached the reconnected client") + } + + // the reconnected handler's own cleanup still removes the registration + ms.deleteClient(clientName, newChan) + ms.clientChansLock.RLock() + _, found := ms.clientChans[clientName] + ms.clientChansLock.RUnlock() + assert.False(t, found) +} + // TestBroadcastVolumeLocationsToClients verifies grown volume locations are sent to registered clients. func TestBroadcastVolumeLocationsToClients(t *testing.T) { clientChan := make(chan *master_pb.KeepConnectedResponse, 2)