Tweaks to pinecone demo to shutdown more cleanly

This commit is contained in:
Devon Hudson 2023-02-24 15:41:47 -07:00
parent 3d31b131fc
commit b28406c7d0
No known key found for this signature in database
GPG Key ID: CD06B18E77F6A628
3 changed files with 62 additions and 28 deletions

View File

@ -72,8 +72,10 @@ type P2PMonolith struct {
port int
httpMux *mux.Router
pineconeMux *mux.Router
httpServer *http.Server
listener net.Listener
httpListenAddr string
stopHandlingEvents chan bool
}
func GenerateDefaultConfig(sk ed25519.PrivateKey, storageDir string, cacheDir string, dbPrefix string) *config.Dendrite {
@ -199,8 +201,10 @@ func (p *P2PMonolith) StartMonolith() {
}
func (p *P2PMonolith) Stop() {
logrus.Info("Stopping monolith")
_ = p.BaseDendrite.Close()
p.WaitForShutdown()
logrus.Info("Stopped monolith")
}
func (p *P2PMonolith) WaitForShutdown() {
@ -209,6 +213,16 @@ func (p *P2PMonolith) WaitForShutdown() {
}
func (p *P2PMonolith) closeAllResources() {
logrus.Info("Closing monolith resources")
if p.httpServer != nil {
p.httpServer.Shutdown(context.Background())
}
select {
case p.stopHandlingEvents <- true:
default:
}
if p.listener != nil {
_ = p.listener.Close()
}
@ -224,6 +238,7 @@ func (p *P2PMonolith) closeAllResources() {
if p.Router != nil {
_ = p.Router.Close()
}
logrus.Info("Monolith resources closed")
}
func (p *P2PMonolith) Addr() string {
@ -280,7 +295,7 @@ func (p *P2PMonolith) setupHttpServers(userProvider *users.PineconeUserProvider,
func (p *P2PMonolith) startHTTPServers() {
go func() {
// Build both ends of a HTTP multiplex.
httpServer := &http.Server{
p.httpServer = &http.Server{
Addr: ":0",
TLSNextProto: map[string]func(*http.Server, *tls.Conn, http.Handler){},
ReadTimeout: 10 * time.Second,
@ -296,12 +311,13 @@ func (p *P2PMonolith) startHTTPServers() {
pubkeyString := hex.EncodeToString(pubkey[:])
logrus.Info("Listening on ", pubkeyString)
switch httpServer.Serve(p.Sessions.Protocol(SessionProtocol)) {
switch p.httpServer.Serve(p.Sessions.Protocol(SessionProtocol)) {
case net.ErrClosed, http.ErrServerClosed:
logrus.Info("Stopped listening on ", pubkeyString)
default:
logrus.Error("Stopped listening on ", pubkeyString)
}
logrus.Info("Stopped goroutine listening on ", pubkeyString)
}()
p.httpListenAddr = fmt.Sprintf(":%d", p.port)
@ -313,10 +329,12 @@ func (p *P2PMonolith) startHTTPServers() {
default:
logrus.Error("Stopped listening on ", p.httpListenAddr)
}
logrus.Info("Stopped goroutine listening on ", p.httpListenAddr)
}()
}
func (p *P2PMonolith) startEventHandler() {
p.stopHandlingEvents = make(chan bool)
stopRelayServerSync := make(chan bool)
eLog := logrus.WithField("pinecone", "events")
p.RelayRetriever = relay.NewRelayServerRetriever(
@ -329,13 +347,19 @@ func (p *P2PMonolith) startEventHandler() {
p.RelayRetriever.InitializeRelayServers(eLog)
go func(ch <-chan pineconeEvents.Event) {
for event := range ch {
for {
select {
case event := <-ch:
switch e := event.(type) {
case pineconeEvents.PeerAdded:
p.RelayRetriever.StartSync()
case pineconeEvents.PeerRemoved:
if p.RelayRetriever.IsRunning() && p.Router.TotalPeerCount() == 0 {
stopRelayServerSync <- true
// NOTE: Don't block on channel
select {
case stopRelayServerSync <- true:
default:
}
}
case pineconeEvents.BroadcastReceived:
// eLog.Info("Broadcast received from: ", e.PeerID)
@ -347,8 +371,17 @@ func (p *P2PMonolith) startEventHandler() {
if err := p.dendrite.FederationAPI.PerformWakeupServers(p.BaseDendrite.Context(), req, res); err != nil {
eLog.WithError(err).Error("Failed to wakeup destination", e.PeerID)
}
}
case <-p.stopHandlingEvents:
logrus.Info("Stopping processing pinecone events")
// NOTE: Don't block on channel
select {
case stopRelayServerSync <- true:
default:
}
logrus.Info("Stopped processing pinecone events")
return
}
}
}(p.EventChannel)
}

View File

@ -38,7 +38,7 @@ type RelayServerRetriever struct {
relayServersQueried map[gomatrixserverlib.ServerName]bool
queriedServersMutex sync.Mutex
running atomic.Bool
quit <-chan bool
quit chan bool
}
func NewRelayServerRetriever(
@ -46,7 +46,7 @@ func NewRelayServerRetriever(
serverName gomatrixserverlib.ServerName,
federationAPI federationAPI.FederationInternalAPI,
relayAPI relayServerAPI.RelayInternalAPI,
quit <-chan bool,
quit chan bool,
) RelayServerRetriever {
return RelayServerRetriever{
ctx: ctx,
@ -151,6 +151,7 @@ func (r *RelayServerRetriever) SyncRelayServers(stop <-chan bool) {
if !t.Stop() {
<-t.C
}
logrus.Info("Stopped relay server retriever")
return
case <-t.C:
}

View File

@ -60,7 +60,7 @@ func TestRelayRetrieverInitialization(t *testing.T) {
"server",
&FakeFedAPI{},
&FakeRelayAPI{},
make(<-chan bool),
make(chan bool),
)
retriever.InitializeRelayServers(logrus.WithField("test", "relay"))
@ -74,7 +74,7 @@ func TestRelayRetrieverSync(t *testing.T) {
"server",
&FakeFedAPI{},
&FakeRelayAPI{},
make(<-chan bool),
make(chan bool),
)
retriever.InitializeRelayServers(logrus.WithField("test", "relay"))