2022-09-05 03:06:08 -04:00
|
|
|
/*
|
|
|
|
Copyright (c) Edgeless Systems GmbH
|
|
|
|
|
|
|
|
SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
*/
|
|
|
|
|
2022-07-14 09:45:04 -04:00
|
|
|
package clean
|
|
|
|
|
|
|
|
import (
|
|
|
|
"sync"
|
|
|
|
)
|
|
|
|
|
2022-11-09 09:57:54 -05:00
|
|
|
// Cleaner can be used to stop a list of services gracefully.
|
|
|
|
// To register an arbitrary amount of stoppers either use New or With.
|
|
|
|
// Start needs to be called to ready the Cleaner, then Clean will activate it.
|
|
|
|
// Done can be used to wait for Cleaner to run all registered stoppers.
|
2022-10-05 09:02:46 -04:00
|
|
|
type Cleaner struct {
|
2022-07-14 09:45:04 -04:00
|
|
|
stoppers []stopper
|
|
|
|
stopC chan struct{}
|
|
|
|
startOnce sync.Once
|
|
|
|
wg sync.WaitGroup
|
|
|
|
}
|
|
|
|
|
|
|
|
// New creates a new cleaner.
|
2022-10-05 09:02:46 -04:00
|
|
|
func New(stoppers ...stopper) *Cleaner {
|
|
|
|
res := &Cleaner{
|
2022-07-14 09:45:04 -04:00
|
|
|
stoppers: stoppers,
|
|
|
|
stopC: make(chan struct{}, 1),
|
|
|
|
}
|
|
|
|
res.wg.Add(1) // for the Start goroutine
|
|
|
|
return res
|
|
|
|
}
|
|
|
|
|
|
|
|
// With adds a new stopper to the cleaner.
|
2022-10-05 09:02:46 -04:00
|
|
|
func (c *Cleaner) With(stopper stopper) *Cleaner {
|
2022-07-14 09:45:04 -04:00
|
|
|
c.stoppers = append(c.stoppers, stopper)
|
|
|
|
return c
|
|
|
|
}
|
|
|
|
|
|
|
|
// Start blocks until it receives a stop message, stops all services gracefully and returns.
|
2022-10-05 09:02:46 -04:00
|
|
|
func (c *Cleaner) Start() {
|
2022-07-14 09:45:04 -04:00
|
|
|
c.startOnce.Do(func() {
|
|
|
|
defer c.wg.Done()
|
|
|
|
// wait for the stop message
|
|
|
|
<-c.stopC
|
|
|
|
|
|
|
|
c.wg.Add(len(c.stoppers))
|
|
|
|
for _, stopItem := range c.stoppers {
|
|
|
|
go func(stopItem stopper) {
|
|
|
|
defer c.wg.Done()
|
|
|
|
stopItem.Stop()
|
|
|
|
}(stopItem)
|
|
|
|
}
|
|
|
|
})
|
|
|
|
}
|
|
|
|
|
|
|
|
// Clean initiates the cleanup but does not wait for it to complete.
|
2022-10-05 09:02:46 -04:00
|
|
|
func (c *Cleaner) Clean() {
|
2022-07-14 09:45:04 -04:00
|
|
|
// try to enqueue the stop message once
|
|
|
|
// if the channel is full, the message is dropped
|
|
|
|
select {
|
|
|
|
case c.stopC <- struct{}{}:
|
|
|
|
default:
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// Done waits for the cleanup to complete.
|
2022-10-05 09:02:46 -04:00
|
|
|
func (c *Cleaner) Done() {
|
2022-07-14 09:45:04 -04:00
|
|
|
c.wg.Wait()
|
|
|
|
}
|
|
|
|
|
|
|
|
type stopper interface {
|
|
|
|
Stop()
|
|
|
|
}
|