main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/0xef53/nsq-consumer"
	"github.com/nsqio/go-nsq"
)

var (
	configFile   = "./foreman.ini"
	logLevelName = "WARN"
	logLevel     = nsq.LogLevelWarning
)

// Consumers is a pool of consumers.
type Consumers map[string]*consumer.Consumer

// Append addss a new consumer to the pool and starts it.
func (c *Consumers) Append(id, name string, t *TopicParams) error {
	(*c)[name] = consumer.NewConsumer(name, t.Channel)

	(*c)[name].SetLogger(Logger, logLevel)

	if len(id) != 0 {
		(*c)[name].Set("client_id", id)
	}

	(*c)[name].Set("nsqlookupd", t.Servers)
	(*c)[name].Set("concurrency", t.Concurrency)
	(*c)[name].Set("max_attempts", t.MaxAttempts)
	(*c)[name].Set("max_in_flight", t.Concurrency)
	(*c)[name].Set("default_requeue_delay", "15s")

	(*c)[name].Start(nsq.HandlerFunc(func(msg *nsq.Message) error {
		data := make(map[string]interface{})
		if err := json.Unmarshal(msg.Body, &data); err != nil {
			return err
		}

		data["_topic_name"] = name
		data["_worker_exit_code"] = 0
		data["_notify_start_exit_code"] = 0

		// It's a prefix preceding each line generated by scripts of this topic
		var prefix string
		if taskID, ok := data["task_id"]; ok {
			prefix = fmt.Sprintf("[%s, %s, %s]", name, msg.ID, taskID)
		} else {
			prefix = fmt.Sprintf("[%s, %s]", name, msg.ID)
		}

		Logger.Printf("%s Got a new message from %s (attempt = %d)\n", prefix, msg.NSQDAddress, msg.Attempts)

		// Notification of the start
		if len(t.NotifyStart) > 0 && msg.Attempts == 1 {
			Logger.Println(prefix, "notify-start command:", t.NotifyStart)
			if exitCode, err := executeCommand(t.NotifyStart, prefix, data, msg, t.Envs); err != nil {
				data["_notify_start_exit_code"] = exitCode
				Logger.Println(prefix, "notify-start error:", err)
			}
		}

		// Executing a main worker command
		switch exitCode, err := executeCommand(t.Cmd, prefix, data, msg, t.Envs); {
		case err == nil:
			// Notification of successful completion
			if len(t.NotifyFinish) > 0 {
				Logger.Println(prefix, "notify-finish command:", t.NotifyFinish)
				if _, err := executeCommand(t.NotifyFinish, prefix, data, msg, t.Envs); err != nil {
					Logger.Println(prefix, "notify-finish error:", err)
				}
			}
		default:
			Logger.Println(prefix, "worker error:", err)
			if IsRequeueError(err) && int(msg.Attempts) < t.MaxAttempts {
				return err
			}
			data["_worker_exit_code"] = exitCode
			// Notification of fault completion
			if len(t.NotifyFault) > 0 {
				Logger.Println(prefix, "notify-fault command:", t.NotifyFault)
				if _, err := executeCommand(t.NotifyFault, prefix, data, msg, t.Envs); err != nil {
					Logger.Println(prefix, "notify-fault error:", err)
				}
			}
		}

		return nil
	}))

	return nil
}

func main() {
	flag.StringVar(&configFile, "c", configFile, "configuration `file`")
	flag.StringVar(&logLevelName, "log-level", logLevelName, "verbosity `level` that is used when logging messages")
	flag.BoolVar(&ShowVersion, "version", false, "print version information and quit")
	flag.Parse()

	if ShowVersion {
		var s string
		if len(CommitRevision) > 0 {
			s = fmt.Sprintf("v%s-%s (built w/%s)", Version, CommitRevision, runtime.Version())
		} else {
			s = fmt.Sprintf("v%s (built w/%s)", Version, runtime.Version())
		}
		fmt.Println(s)
		return
	}

	Logger.Println("Using config file", configFile)

	cfg, err := NewConfig(configFile)
	if err != nil {
		Logger.Fatalln(err)
	}

	switch logLevelName {
	case "DEBUG":
		logLevel = nsq.LogLevelDebug
	case "INFO":
		logLevel = nsq.LogLevelInfo
	case "WARN":
		logLevel = nsq.LogLevelWarning
	case "ERROR":
		logLevel = nsq.LogLevelError
	default:
		Logger.Fatalln("error: unknown log level:", logLevelName)
	}

	pool := make(Consumers)

	for name, opts := range cfg.Topic {
		if err := pool.Append(cfg.Common.ClientID, name, opts); err != nil {
			Logger.Fatalln(err)
		}
	}

	cNowExit := make(chan struct{})

	cSignal := make(chan os.Signal, 1)
	signal.Notify(cSignal, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	go func() {
		s := <-cSignal
		Logger.Printf("Captured signal: %s. Gracefully stopping all consumers and exiting...\n", s)
		for name, c := range pool {
			go func(name string, c *consumer.Consumer) {
				wg.Add(1)
				defer wg.Done()
				Logger.Printf("Closing %s...\n", name)
				c.Stop()
				Logger.Printf("Successfully closed %s\n", name)
			}(name, c)
		}
		close(cNowExit)
	}()

	<-cNowExit

	// Waiting for all goroutines
	wg.Wait()

	Logger.Println("All consumers has been successfully stopped. Exiting with a clear conscience.")

	os.Exit(0)
}