initial import

core/proc/env.go

@@ -0,0 +1,43 @@
+package proc
+
+import (
+	"os"
+	"strconv"
+	"sync"
+)
+
+var (
+	envs    = make(map[string]string)
+	envLock sync.RWMutex
+)
+
+func Env(name string) string {
+	envLock.RLock()
+	val, ok := envs[name]
+	envLock.RUnlock()
+
+	if ok {
+		return val
+	}
+
+	val = os.Getenv(name)
+	envLock.Lock()
+	envs[name] = val
+	envLock.Unlock()
+
+	return val
+}
+
+func EnvInt(name string) (int, bool) {
+	val := Env(name)
+	if len(val) == 0 {
+		return 0, false
+	}
+
+	n, err := strconv.Atoi(val)
+	if err != nil {
+		return 0, false
+	}
+
+	return n, true
+}

core/proc/env_test.go

@@ -0,0 +1,34 @@
+package proc
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEnv(t *testing.T) {
+	assert.True(t, len(Env("any")) == 0)
+	envLock.RLock()
+	val, ok := envs["any"]
+	envLock.RUnlock()
+	assert.True(t, len(val) == 0)
+	assert.True(t, ok)
+	assert.True(t, len(Env("any")) == 0)
+}
+
+func TestEnvInt(t *testing.T) {
+	val, ok := EnvInt("any")
+	assert.Equal(t, 0, val)
+	assert.False(t, ok)
+	err := os.Setenv("anyInt", "10")
+	assert.Nil(t, err)
+	val, ok = EnvInt("anyInt")
+	assert.Equal(t, 10, val)
+	assert.True(t, ok)
+	err = os.Setenv("anyString", "a")
+	assert.Nil(t, err)
+	val, ok = EnvInt("anyString")
+	assert.Equal(t, 0, val)
+	assert.False(t, ok)
+}

core/proc/goroutines+polyfill.go

@@ -0,0 +1,6 @@
+// +build windows
+
+package proc
+
+func dumpGoroutines() {
+}

core/proc/goroutines.go

@@ -0,0 +1,35 @@
+// +build linux darwin
+
+package proc
+
+import (
+	"fmt"
+	"os"
+	"path"
+	"runtime/pprof"
+	"syscall"
+	"time"
+
+	"zero/core/logx"
+)
+
+const (
+	goroutineProfile = "goroutine"
+	debugLevel       = 2
+)
+
+func dumpGoroutines() {
+	command := path.Base(os.Args[0])
+	pid := syscall.Getpid()
+	dumpFile := path.Join(os.TempDir(), fmt.Sprintf("%s-%d-goroutines-%s.dump",
+		command, pid, time.Now().Format(timeFormat)))
+
+	logx.Infof("Got dump goroutine signal, printing goroutine profile to %s", dumpFile)
+
+	if f, err := os.Create(dumpFile); err != nil {
+		logx.Errorf("Failed to dump goroutine profile, error: %v", err)
+	} else {
+		defer f.Close()
+		pprof.Lookup(goroutineProfile).WriteTo(f, debugLevel)
+	}
+}

core/proc/process.go

@@ -0,0 +1,24 @@
+package proc
+
+import (
+	"os"
+	"path/filepath"
+)
+
+var (
+	procName string
+	pid      int
+)
+
+func init() {
+	procName = filepath.Base(os.Args[0])
+	pid = os.Getpid()
+}
+
+func Pid() int {
+	return pid
+}
+
+func ProcessName() string {
+	return procName
+}

core/proc/process_test.go

@@ -0,0 +1,15 @@
+package proc
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestProcessName(t *testing.T) {
+	assert.True(t, len(ProcessName()) > 0)
+}
+
+func TestPid(t *testing.T) {
+	assert.True(t, Pid() > 0)
+}

core/proc/profile+polyfill.go

@@ -0,0 +1,7 @@
+// +build windows
+
+package proc
+
+func StartProfile() Stopper {
+	return noopStopper
+}

core/proc/profile.go

@@ -0,0 +1,205 @@
+// +build linux darwin
+
+package proc
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"path"
+	"runtime"
+	"runtime/pprof"
+	"runtime/trace"
+	"sync/atomic"
+	"syscall"
+	"time"
+
+	"zero/core/logx"
+)
+
+// DefaultMemProfileRate is the default memory profiling rate.
+// See also http://golang.org/pkg/runtime/#pkg-variables
+const DefaultMemProfileRate = 4096
+
+// started is non zero if a profile is running.
+var started uint32
+
+// Profile represents an active profiling session.
+type Profile struct {
+	// path holds the base path where various profiling files are  written.
+	// If blank, the base path will be generated by ioutil.TempDir.
+	path string
+
+	// closers holds cleanup functions that run after each profile
+	closers []func()
+
+	// stopped records if a call to profile.Stop has been made
+	stopped uint32
+}
+
+func (p *Profile) close() {
+	for _, closer := range p.closers {
+		closer()
+	}
+}
+
+func (p *Profile) startBlockProfile() {
+	fn := createDumpFile("block")
+	f, err := os.Create(fn)
+	if err != nil {
+		logx.Errorf("profile: could not create block profile %q: %v", fn, err)
+		return
+	}
+
+	runtime.SetBlockProfileRate(1)
+	logx.Infof("profile: block profiling enabled, %s", fn)
+	p.closers = append(p.closers, func() {
+		pprof.Lookup("block").WriteTo(f, 0)
+		f.Close()
+		runtime.SetBlockProfileRate(0)
+		logx.Infof("profile: block profiling disabled, %s", fn)
+	})
+}
+
+func (p *Profile) startCpuProfile() {
+	fn := createDumpFile("cpu")
+	f, err := os.Create(fn)
+	if err != nil {
+		logx.Errorf("profile: could not create cpu profile %q: %v", fn, err)
+		return
+	}
+
+	logx.Infof("profile: cpu profiling enabled, %s", fn)
+	pprof.StartCPUProfile(f)
+	p.closers = append(p.closers, func() {
+		pprof.StopCPUProfile()
+		f.Close()
+		logx.Infof("profile: cpu profiling disabled, %s", fn)
+	})
+}
+
+func (p *Profile) startMemProfile() {
+	fn := createDumpFile("mem")
+	f, err := os.Create(fn)
+	if err != nil {
+		logx.Errorf("profile: could not create memory profile %q: %v", fn, err)
+		return
+	}
+
+	old := runtime.MemProfileRate
+	runtime.MemProfileRate = DefaultMemProfileRate
+	logx.Infof("profile: memory profiling enabled (rate %d), %s", runtime.MemProfileRate, fn)
+	p.closers = append(p.closers, func() {
+		pprof.Lookup("heap").WriteTo(f, 0)
+		f.Close()
+		runtime.MemProfileRate = old
+		logx.Infof("profile: memory profiling disabled, %s", fn)
+	})
+}
+
+func (p *Profile) startMutexProfile() {
+	fn := createDumpFile("mutex")
+	f, err := os.Create(fn)
+	if err != nil {
+		logx.Errorf("profile: could not create mutex profile %q: %v", fn, err)
+		return
+	}
+
+	runtime.SetMutexProfileFraction(1)
+	logx.Infof("profile: mutex profiling enabled, %s", fn)
+	p.closers = append(p.closers, func() {
+		if mp := pprof.Lookup("mutex"); mp != nil {
+			mp.WriteTo(f, 0)
+		}
+		f.Close()
+		runtime.SetMutexProfileFraction(0)
+		logx.Infof("profile: mutex profiling disabled, %s", fn)
+	})
+}
+
+func (p *Profile) startThreadCreateProfile() {
+	fn := createDumpFile("threadcreate")
+	f, err := os.Create(fn)
+	if err != nil {
+		logx.Errorf("profile: could not create threadcreate profile %q: %v", fn, err)
+		return
+	}
+
+	logx.Infof("profile: threadcreate profiling enabled, %s", fn)
+	p.closers = append(p.closers, func() {
+		if mp := pprof.Lookup("threadcreate"); mp != nil {
+			mp.WriteTo(f, 0)
+		}
+		f.Close()
+		logx.Infof("profile: threadcreate profiling disabled, %s", fn)
+	})
+}
+
+func (p *Profile) startTraceProfile() {
+	fn := createDumpFile("trace")
+	f, err := os.Create(fn)
+	if err != nil {
+		logx.Errorf("profile: could not create trace output file %q: %v", fn, err)
+		return
+	}
+
+	if err := trace.Start(f); err != nil {
+		logx.Errorf("profile: could not start trace: %v", err)
+		return
+	}
+
+	logx.Infof("profile: trace enabled, %s", fn)
+	p.closers = append(p.closers, func() {
+		trace.Stop()
+		logx.Infof("profile: trace disabled, %s", fn)
+	})
+}
+
+// Stop stops the profile and flushes any unwritten data.
+func (p *Profile) Stop() {
+	if !atomic.CompareAndSwapUint32(&p.stopped, 0, 1) {
+		// someone has already called close
+		return
+	}
+	p.close()
+	atomic.StoreUint32(&started, 0)
+}
+
+// Start starts a new profiling session.
+// The caller should call the Stop method on the value returned
+// to cleanly stop profiling.
+func StartProfile() Stopper {
+	if !atomic.CompareAndSwapUint32(&started, 0, 1) {
+		logx.Error("profile: Start() already called")
+		return noopStopper
+	}
+
+	var prof Profile
+	prof.startCpuProfile()
+	prof.startMemProfile()
+	prof.startMutexProfile()
+	prof.startBlockProfile()
+	prof.startTraceProfile()
+	prof.startThreadCreateProfile()
+
+	go func() {
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, syscall.SIGINT)
+		<-c
+
+		logx.Info("profile: caught interrupt, stopping profiles")
+		prof.Stop()
+
+		signal.Reset()
+		syscall.Kill(os.Getpid(), syscall.SIGINT)
+	}()
+
+	return &prof
+}
+
+func createDumpFile(kind string) string {
+	command := path.Base(os.Args[0])
+	pid := syscall.Getpid()
+	return path.Join(os.TempDir(), fmt.Sprintf("%s-%d-%s-%s.pprof",
+		command, pid, kind, time.Now().Format(timeFormat)))
+}

core/proc/shutdown+polyfill.go

@@ -0,0 +1,16 @@
+// +build windows
+
+package proc
+
+import "time"
+
+func AddShutdownListener(fn func()) func() {
+	return nil
+}
+
+func AddWrapUpListener(fn func()) func() {
+	return nil
+}
+
+func SetTimeoutToForceQuit(duration time.Duration) {
+}

core/proc/shutdown.go

@@ -0,0 +1,81 @@
+// +build linux darwin
+
+package proc
+
+import (
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+	"time"
+
+	"zero/core/logx"
+)
+
+const (
+	wrapUpTime = time.Second
+	// why we use 5500 milliseconds is because most of our queue are blocking mode with 5 seconds
+	waitTime = 5500 * time.Millisecond
+)
+
+var (
+	wrapUpListeners          = new(listenerManager)
+	shutdownListeners        = new(listenerManager)
+	delayTimeBeforeForceQuit = waitTime
+)
+
+func AddShutdownListener(fn func()) (waitForCalled func()) {
+	return shutdownListeners.addListener(fn)
+}
+
+func AddWrapUpListener(fn func()) (waitForCalled func()) {
+	return wrapUpListeners.addListener(fn)
+}
+
+func SetTimeoutToForceQuit(duration time.Duration) {
+	delayTimeBeforeForceQuit = duration
+}
+
+func gracefulStop(signals chan os.Signal) {
+	signal.Stop(signals)
+
+	logx.Info("Got signal SIGTERM, shutting down...")
+	wrapUpListeners.notifyListeners()
+
+	time.Sleep(wrapUpTime)
+	shutdownListeners.notifyListeners()
+
+	time.Sleep(delayTimeBeforeForceQuit - wrapUpTime)
+	logx.Infof("Still alive after %v, going to force kill the process...", delayTimeBeforeForceQuit)
+	syscall.Kill(syscall.Getpid(), syscall.SIGTERM)
+}
+
+type listenerManager struct {
+	lock      sync.Mutex
+	waitGroup sync.WaitGroup
+	listeners []func()
+}
+
+func (lm *listenerManager) addListener(fn func()) (waitForCalled func()) {
+	lm.waitGroup.Add(1)
+
+	lm.lock.Lock()
+	lm.listeners = append(lm.listeners, func() {
+		defer lm.waitGroup.Done()
+		fn()
+	})
+	lm.lock.Unlock()
+
+	return func() {
+		lm.waitGroup.Wait()
+	}
+}
+
+func (lm *listenerManager) notifyListeners() {
+	lm.lock.Lock()
+	defer lm.lock.Unlock()
+
+	for _, listener := range lm.listeners {
+		listener()
+	}
+}

core/proc/signals.go

@@ -0,0 +1,42 @@
+// +build linux darwin
+
+package proc
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+
+	"zero/core/logx"
+)
+
+const timeFormat = "0102150405"
+
+func init() {
+	go func() {
+		var profiler Stopper
+
+		// https://golang.org/pkg/os/signal/#Notify
+		signals := make(chan os.Signal, 1)
+		signal.Notify(signals, syscall.SIGUSR1, syscall.SIGUSR2, syscall.SIGTERM)
+
+		for {
+			v := <-signals
+			switch v {
+			case syscall.SIGUSR1:
+				dumpGoroutines()
+			case syscall.SIGUSR2:
+				if profiler == nil {
+					profiler = StartProfile()
+				} else {
+					profiler.Stop()
+					profiler = nil
+				}
+			case syscall.SIGTERM:
+				gracefulStop(signals)
+			default:
+				logx.Error("Got unregistered signal:", v)
+			}
+		}
+	}()
+}

core/proc/stopper.go

@@ -0,0 +1,14 @@
+package proc
+
+var noopStopper nilStopper
+
+type (
+	Stopper interface {
+		Stop()
+	}
+
+	nilStopper struct{}
+)
+
+func (ns nilStopper) Stop() {
+}