Async Hook added

markdicksonjr · Dec 19, 2017 · bccae70 · bccae70
1 parent 6d12892
commit bccae70
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 8 deletions.
diff --git a/hook.go b/hook.go
@@ -19,6 +19,8 @@ var (
 
 type IndexNameFunc func() string
 
+type fireFunc func(entry *logrus.Entry, hook *ElasticHook) error
+
 // ElasticHook is a logrus
 // hook for ElasticSearch
 type ElasticHook struct {
@@ -28,6 +30,7 @@ type ElasticHook struct {
 	levels    []logrus.Level
 	ctx       context.Context
 	ctxCancel context.CancelFunc
+	fireFunc  fireFunc
 }
 
 // NewElasticHook creates new hook
@@ -39,6 +42,15 @@ func NewElasticHook(client *elastic.Client, host string, level logrus.Level, ind
 	return NewElasticHookWithFunc(client, host, level, func() string { return index })
 }
 
+// NewElasticHook creates new  hook with asynchronous log
+// client - ElasticSearch client using gopkg.in/olivere/elastic.v5
+// host - host of system
+// level - log level
+// index - name of the index in ElasticSearch
+func NewAsyncElasticHook(client *elastic.Client, host string, level logrus.Level, index string) (*ElasticHook, error) {
+	return NewAsyncElasticHookWithFunc(client, host, level, func() string { return index })
+}
+
 // NewElasticHookWithFunc creates new hook with
 // function that provides the index name. This is useful if the index name is
 // somehow dynamic especially based on time.
@@ -47,6 +59,21 @@ func NewElasticHook(client *elastic.Client, host string, level logrus.Level, ind
 // level - log level
 // indexFunc - function providing the name of index
 func NewElasticHookWithFunc(client *elastic.Client, host string, level logrus.Level, indexFunc IndexNameFunc) (*ElasticHook, error) {
+	return newHookFuncAndFireFunc(client, host, level, indexFunc, syncFireFunc)
+}
+
+// NewAsyncElasticHookWithFunc creates new asynchronous hook with
+// function that provides the index name. This is useful if the index name is
+// somehow dynamic especially based on time.
+// client - ElasticSearch client using gopkg.in/olivere/elastic.v5
+// host - host of system
+// level - log level
+// indexFunc - function providing the name of index
+func NewAsyncElasticHookWithFunc(client *elastic.Client, host string, level logrus.Level, indexFunc IndexNameFunc) (*ElasticHook, error) {
+	return newHookFuncAndFireFunc(client, host, level, indexFunc, asyncFireFunc)
+}
+
+func newHookFuncAndFireFunc(client *elastic.Client, host string, level logrus.Level, indexFunc IndexNameFunc, fireFunc fireFunc) (*ElasticHook, error) {
 	levels := []logrus.Level{}
 	for _, l := range []logrus.Level{
 		logrus.PanicLevel,
@@ -86,13 +113,22 @@ func NewElasticHookWithFunc(client *elastic.Client, host string, level logrus.Le
 		levels:    levels,
 		ctx:       ctx,
 		ctxCancel: cancel,
+		fireFunc:  fireFunc,
 	}, nil
 }
 
 // Fire is required to implement
 // Logrus hook
 func (hook *ElasticHook) Fire(entry *logrus.Entry) error {
+	return hook.fireFunc(entry, hook)
+}
+
+func asyncFireFunc(entry *logrus.Entry, hook *ElasticHook) error {
+	go syncFireFunc(entry, hook)
+	return nil
+}
 
+func syncFireFunc(entry *logrus.Entry, hook *ElasticHook) error {
 	level := entry.Level.String()
 
 	if e, ok := entry.Data[logrus.ErrorKey]; ok && e != nil {

diff --git a/hook_test.go b/hook_test.go
@@ -19,13 +19,23 @@ import (
 
 //docker run -it --rm -p 7777:9200 -p 5601:5601 sebp/elk
 
+type NewHookFunc func(client *elastic.Client, host string, level logrus.Level, index string) (*ElasticHook, error)
+
 type Log struct{}
 
 func (l Log) Printf(format string, args ...interface{}) {
 	log.Printf(format+"\n", args)
 }
 
-func TestHook(t *testing.T) {
+func TestSyncHook(t *testing.T) {
+	hookTest(NewElasticHook, t)
+}
+
+func TestAsyncHook(t *testing.T) {
+	hookTest(NewAsyncElasticHook, t)
+}
+
+func hookTest(hookfunc NewHookFunc, t *testing.T) {
 	if r, err := http.Get("http://127.0.0.1:7777"); err != nil {
 		log.Fatal("Elastic not reachable")
 	} else {
@@ -43,6 +53,10 @@ func TestHook(t *testing.T) {
 		log.Panic(err)
 	}
 
+	client.
+		DeleteIndex("goplag").
+		Do(context.TODO())
+
 	hook, err := NewElasticHook(client, "localhost", logrus.DebugLevel, "goplag")
 	if err != nil {
 		log.Panic(err)
@@ -63,7 +77,7 @@ func TestHook(t *testing.T) {
 		Do(context.TODO())
 
 	if searchResult.Hits.TotalHits != 100 {
-		t.Error("Not all logs pushed to elastic")
+		t.Errorf("Not all logs pushed to elastic: expected %d got %d", 100, searchResult.Hits.TotalHits)
 		t.FailNow()
 	}
 }
@@ -78,15 +92,10 @@ func TestError(t *testing.T) {
 		log.Panic(err)
 	}
 
-	_, err = client.
+	client.
 		DeleteIndex("errorlog").
 		Do(context.TODO())
 
-	if err != nil {
-		t.Error(err.Error())
-		t.FailNow()
-	}
-
 	time.Sleep(1 * time.Second)
 
 	hook, err := NewElasticHook(client, "localhost", logrus.DebugLevel, "errorlog")