Graceful shutdown handler for Martini

main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/go-martini/martini"
	"net/http"
	"sync/atomic"
)

type GracefulShutdown struct {
	timeout time.Duration
	wg      sync.WaitGroup
}

func NewGracefulShutdown(t time.Duration) *GracefulShutdown {
	return &GracefulShutdown{timeout: t}
}

func (g *GracefulShutdown) Handler(c martini.Context) {
	g.wg.Add(1)
	c.Next()
	g.wg.Done()
}

func (g *GracefulShutdown) WaitForSignal(signals ...os.Signal) error {
	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, signals...)
	<-sigchan

	log.Println("Waiting for all requests to finish")

	waitChan := make(chan struct{})
	go func() {
		g.wg.Wait()
		waitChan <- struct{}{}
	}()

	select {
	case <-time.After(g.timeout):
		return fmt.Errorf("timed out waiting %v for shutdown", g.timeout)
	case <-waitChan:
		return nil
	}
}

type ConnectionLimit struct {
	numConnections int32
	limit          int32
}

func (c *ConnectionLimit) Handler(ctx martini.Context, rw http.ResponseWriter) {
	if atomic.AddInt32(&c.numConnections, 1) > c.limit {
		http.Error(rw, "maximum connections exceeded", http.StatusServiceUnavailable)
		atomic.AddInt32(&c.numConnections, -1)
		return
	}

	ctx.Next()
	atomic.AddInt32(&c.numConnections, -1)
}

func main() {
	m := martini.Classic()

	timeout := flag.Int("timeout", 10, "seconds to wait for a graceful shutdown")
	flag.Parse()

	gracefulShutdown := &GracefulShutdown{timeout: time.Duration(*timeout) * time.Second}
	m.Use(gracefulShutdown.Handler)

	limiter := &ConnectionLimit{limit: 2}
	m.Use(limiter.Handler)

	m.Get("/", func() string {
		time.Sleep(5 * time.Second)
		return "hello world\n"
	})

	go func() {
		m.Run()
	}()

	err := gracefulShutdown.WaitForSignal(syscall.SIGTERM, syscall.SIGINT)
	if err != nil {
		log.Println(err)
	}
}