From c2ce49234de6a09507a4fd356b82a21e97001bb4 Mon Sep 17 00:00:00 2001 From: Maciej Lisiewski Date: Wed, 7 Jun 2017 15:54:01 -0400 Subject: [PATCH] Add pause/resume logic Allows a proxy to be paused (with optional timeout for already open connections) and resumed (optionally with staggered connection open) as well as for the configuration to be updated for new connections only. All that is preparatory work for graceful switchover support in Stolon. --- pollon.go | 161 ++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 137 insertions(+), 24 deletions(-) diff --git a/pollon.go b/pollon.go index 5f34dcf..54c453b 100644 --- a/pollon.go +++ b/pollon.go @@ -19,6 +19,7 @@ import ( "io" "net" "sync" + "time" ) type ConfData struct { @@ -26,45 +27,60 @@ type ConfData struct { } type Proxy struct { - C chan ConfData - listener *net.TCPListener - confMutex sync.Mutex - destAddr *net.TCPAddr + listener *net.TCPListener + + configMutex sync.Mutex + C chan ConfData + destAddr *net.TCPAddr + + connMutex sync.Mutex closeConns chan struct{} - stop chan struct{} + stopCh chan struct{} endCh chan error - connMutex sync.Mutex + + limiter *time.Ticker + state int } +const ( + proxyActive = iota + proxyPaused +) + func NewProxy(listener *net.TCPListener) (*Proxy, error) { return &Proxy{ - C: make(chan ConfData), - listener: listener, - closeConns: make(chan struct{}), - stop: make(chan struct{}), - endCh: make(chan error), - connMutex: sync.Mutex{}, + listener: listener, + configMutex: sync.Mutex{}, + C: make(chan ConfData), + connMutex: sync.Mutex{}, + closeConns: make(chan struct{}), + stopCh: make(chan struct{}), + endCh: make(chan error), }, nil } func (p *Proxy) proxyConn(conn *net.TCPConn) { p.connMutex.Lock() + if p.limiter != nil { + <-p.limiter.C + } closeConns := p.closeConns destAddr := p.destAddr p.connMutex.Unlock() + defer func() { log.Printf("closing source connection: %v", conn) conn.Close() }() - defer conn.Close() if destAddr == nil { + log.Print("can't open connection: destination address empty") return } destConn, err := net.DialTCP("tcp", nil, p.destAddr) if err != nil { - conn.Close() + log.Printf("failed to open connection to destination: %v", err) return } defer func() { @@ -73,36 +89,39 @@ func (p *Proxy) proxyConn(conn *net.TCPConn) { }() var wg sync.WaitGroup - end := make(chan bool) + end := make(chan struct{}) wg.Add(1) go func() { - defer wg.Done() n, err := io.Copy(destConn, conn) if err != nil { + log.Printf("error copying from source to dest: %v", err) } - conn.Close() + conn.CloseWrite() destConn.CloseRead() log.Printf("ending. copied %d bytes from source to dest", n) + wg.Done() }() + wg.Add(1) go func() { - defer wg.Done() n, err := io.Copy(conn, destConn) if err != nil { + log.Printf("error copying from dest to source: %v", err) } - destConn.Close() + destConn.CloseWrite() conn.CloseRead() log.Printf("ending. copied %d bytes from dest to source", n) + wg.Done() }() go func() { wg.Wait() - end <- true + close(end) }() select { case <-end: - log.Printf("all io copy goroutines done") + log.Printf("connection closing - all copy goroutines done") return case <-closeConns: log.Printf("closing all connections") @@ -113,12 +132,14 @@ func (p *Proxy) proxyConn(conn *net.TCPConn) { func (p *Proxy) confCheck() { for { select { - case <-p.stop: + case <-p.stopCh: return case confData := <-p.C: if confData.DestAddr.String() != p.destAddr.String() { p.connMutex.Lock() - close(p.closeConns) + if p.closeConns != nil { + close(p.closeConns) + } p.closeConns = make(chan struct{}) p.destAddr = confData.DestAddr p.connMutex.Unlock() @@ -146,9 +167,101 @@ func (p *Proxy) Start() error { go p.confCheck() go p.accepter() err := <-p.endCh - close(p.stop) + close(p.stopCh) if err != nil { return fmt.Errorf("proxy error: %v", err) } return nil } + +// Pause makes proxy continue accepting new connections, but prevents it +// from opening new connections to destination address until proxy status +// changes. +// It doesn't affect already open connections in any way. +func (p *Proxy) Pause() { + p.configMutex.Lock() + if p.state == proxyActive { + p.state = proxyPaused + p.connMutex.Unlock() + } + p.configMutex.Unlock() +} + +// Resume makes a paused proxy resume opening news connections to destination +// address. +// It makes no attempt to spread backlogged connection opens over time, so it +// can cause thundering herd problem - use with caution +func (p *Proxy) Resume() { + p.configMutex.Lock() + if p.state == proxyPaused { + p.state = proxyActive + if p.closeConns == nil { + p.closeConns = make(chan struct{}) + } + p.connMutex.Unlock() + } + p.configMutex.Unlock() +} + +// PauseAndDisconnect first pauses the proxy and then, after timeout t has +// passed, closes all connections opened prior to the proxy being paused. +// It can be used for a graceful switchover when paired with client-side +// connection timeout +func (p *Proxy) PauseAndDisconnect(t time.Duration) { + p.configMutex.Lock() + + if p.state == proxyActive { + p.state = proxyPaused + p.connMutex.Lock() + + closeConns := p.closeConns + p.closeConns = nil + go func() { + <-time.NewTimer(t).C + close(closeConns) + }() + } + + p.configMutex.Unlock() +} + +// GradualResume makes the proxy resume opening connections to destination, +// but to avoid thundering herd problem by opening at the rate of one +// connection every interval i for the period of duration d +func (p *Proxy) GradualResume(i, d time.Duration) { + p.configMutex.Lock() + if p.state == proxyPaused { + p.state = proxyActive + if p.closeConns == nil { + p.closeConns = make(chan struct{}) + } + p.limit(i, d) + p.connMutex.Unlock() + } + p.configMutex.Unlock() +} + +func (p *Proxy) limit(i, d time.Duration) { + if i.Nanoseconds() <= 0 || d.Nanoseconds() <= 0 { + return + } + p.limiter = time.NewTicker(i) + go func() { + <-time.NewTimer(d).C + p.connMutex.Lock() + p.limiter = nil + p.connMutex.Unlock() + }() +} + +// Config atomically updates proxy destination address without affecting +// already open connections +func (p *Proxy) Config(conf ConfData) { + p.configMutex.Lock() + if conf.DestAddr.String() != p.destAddr.String() { + p.connMutex.Lock() + p.destAddr = conf.DestAddr + p.connMutex.Unlock() + } + p.configMutex.Unlock() +}