Heka采集系统Output插件开发

2019独角兽企业重金招聘Python工程师标准>>>

1、插件接口实现：

    注册插件实例    

func init() {
	pipeline.RegisterPlugin("HttpOutput", func() interface{} {
		return new(HttpOutput)
	})
}

    内部数据初始化

func (o *HttpOutput) ConfigStruct() interface{} {
	return &HttpOutputConfig{
		HttpTimeout: 0,
		Headers:     make(http.Header),
		Method:      "POST",
	}
}

func (o *HttpOutput) Init(config interface{}) (err error) {
	o.HttpOutputConfig = config.(*HttpOutputConfig)
	if o.url, err = url.Parse(o.Address); err != nil {
		return fmt.Errorf("Can't parse URL '%s': %s", o.Address, err.Error())
	}
	if o.url.Scheme != "http" && o.url.Scheme != "https" {
		return errors.New("`address` must contain an absolute http or https URL.")
	}
	o.Method = strings.ToUpper(o.Method)
	if o.Method != "POST" && o.Method != "GET" && o.Method != "PUT" {
		return errors.New("HTTP Method must be POST, GET, or PUT.")
	}
	if o.Method != "GET" {
		o.sendBody = true
	}
	o.client = new(http.Client)
	if o.HttpTimeout > 0 {
		o.client.Timeout = time.Duration(o.HttpTimeout) * time.Millisecond
	}
	if o.Username != "" || o.Password != "" {
		o.useBasicAuth = true
	}
	if o.url.Scheme == "https" {
		transport := &http.Transport{}
		if transport.TLSClientConfig, err = tcp.CreateGoTlsConfig(&o.Tls); err != nil {
			return fmt.Errorf("TLS init error: %s", err.Error())
		}
		o.client.Transport = transport
	}
	return
}

    实例运行

func (o *HttpOutput) Run(or pipeline.OutputRunner, h pipeline.PluginHelper) (err error) {
	if or.Encoder() == nil {
		return errors.New("Encoder must be specified.")
	}

	var (
		e        error
		outBytes []byte
	)
	inChan := or.InChan()

	for pack := range inChan {
		// tob, _ := json.Marshal(pack.Message)
		// fmt.Println(string(tob))
		outBytes, e = or.Encode(pack)
		pack.Recycle()
		if e != nil {
			or.LogError(e)
			continue
		}
		if outBytes == nil {
			continue
		}
		if e = o.request(or, outBytes); e != nil {
			or.LogError(e)
		}
	}

	return
}

    一般实例Run接口中的实现基本类似，获取or.InChan，循环读取，调用Encode，回收pack，然后将encode的数据发送到输出源；

2、以Mongodb作为输出源为例，代码如下：

package mongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/mozilla-services/heka/pipeline"
	"labix.org/v2/mgo"
	"log"
	"runtime"
	"runtime/debug"
	"strings"
	"time"
)

type MongoOutput struct {
	*MongoOutputConfig
	mgoSession *mgo.Session
	logMsgChan chan *logMessage
}

type ServerId map[string]string

type MongoOutputConfig struct {
	Address        []string
	LogMsgChSize   int
	MgoWriterCount int
	SrvIds         ServerId
	Username       string `toml:"username"`
	Password       string `toml:"password"`
}

type logMessage struct {
	ServerId  string `json:"Hostname"`
	ClientIp  string `json:"remote_addr"`
	Host      string `json:"host"`
	Cookie    string `json:"http_cookie"`
	Referer   string `json:"http_referer"`
	URI       string `json:"request_uri"`
	Timestamp int64  `json:"timestamp"`
	UserAgent string `json:"http_user_agent"`
}

//data transferred from agent.
type mgoMessage struct {
	Url       string `json:"url"`
	UA        string `json:"ua"`
	Referer   string `json:"referer"`
	Cookie    string `json:"cookie"`
	Ip        string `json:"ip"`
	TimeStamp int64  `json:"timestamp"`
}

func (o *MongoOutput) ConfigStruct() interface{} {
	return &MongoOutputConfig{
		Address:        []string{"127.0.0.1:27017"},
		LogMsgChSize:   10000,
		MgoWriterCount: runtime.NumCPU(),
		SrvIds:         make(ServerId),
	}
}

func (o *MongoOutput) Init(config interface{}) (err error) {
	o.MongoOutputConfig = config.(*MongoOutputConfig)
	//todo: check address validity
	// if o.url, err = url.Parse(o.Address); err != nil {
	// 	return fmt.Errorf("Can't parse URL '%s': %s", o.Address, err.Error())
	// }
	//

	o.logMsgChan = make(chan *logMessage, o.LogMsgChSize)

	mgoInfo := mgo.DialInfo{Addrs: o.Address, Timeout: time.Second}
	o.mgoSession, err = mgo.DialWithInfo(&mgoInfo)
	if err != nil {
		log.Printf("initialize MongoOutput failed, %s", err.Error())
		return err
	}
	return
}

func WriteDataToMgo(mo *MongoOutput) {
	defer func() {
		if err, ok := recover().(error); ok {
			log.Println("WARN: panic in %v", err)
			log.Println(string(debug.Stack()))
		}
	}()
	//srvlog.Printf("WriteDataToMgo key:%s", key)
	sessionCopy := mo.mgoSession.Copy()
	defer sessionCopy.Close()
	db := sessionCopy.DB("admin")
	var msg mgoMessage
	var err error
	for logMsg := range mo.logMsgChan {
		if logMsg.ClientIp == "" || logMsg.UserAgent == "" {
			continue
		}

		keyName := mo.SrvIds[logMsg.ServerId]
		//fmt.Printf("%s, %s", logMsg.ServerId, keyName)

		if keyName == "" {
			log.Printf("no invalid mongo key for %s", logMsg.ServerId)
			continue
		}

		year, month, day := time.Now().Date()
		collName := fmt.Sprintf("%s_%d_%d_%d", keyName, year, month, day)
		coll := db.C(collName)
		msg.Url = fmt.Sprintf("http://%s%s", logMsg.Host, logMsg.URI)
		if logMsg.UserAgent != "-" {
			msg.UA = logMsg.UserAgent
		} else {
			msg.UA = ""
		}

		if logMsg.Referer != "-" {
			msg.Referer = logMsg.Referer
		} else {
			msg.Referer = ""
		}

		if logMsg.Cookie != "-" {
			msg.Cookie = logMsg.Cookie
		} else {
			msg.Cookie = ""
		}

		msg.Ip = logMsg.ClientIp
		msg.TimeStamp = logMsg.Timestamp
		if err != nil {
			log.Println(err.Error())
			continue
		}
		//fmt.Println("MongoOutput-119-%v", msg)
		err = coll.Insert(msg)
		if err != nil {
			log.Println(err.Error())
			continue
		}
	}
}

func (o *MongoOutput) Run(or pipeline.OutputRunner, h pipeline.PluginHelper) (err error) {
	if or.Encoder() == nil {
		return errors.New("Encoder must be specified.")
	}

	var (
		e        error
		outBytes []byte
	)
	inChan := or.InChan()

	for i := 0; i < o.MgoWriterCount; i++ {
		go WriteDataToMgo(o)
	}

	for pack := range inChan {
		outBytes, e = or.Encode(pack)
		pack.Recycle()
		if e != nil {
			or.LogError(e)
			continue
		}
		if outBytes == nil {
			continue
		}
		subStrs := strings.Split(string(outBytes), "\n")
		//fmt.Printf("MongoOutput:len of subStrs:%d", len(subStrs))

		if len(subStrs) == 3 {
			logMsg := &logMessage{}
			//fmt.Printf("MongoOutput:%s\n", subStrs[1])
			e = json.Unmarshal([]byte(subStrs[1]), logMsg)
			if e != nil {
				log.Printf("MongoOutput-%s", err.Error())
				continue
			}

			//fmt.Printf("MongoOutput-165-%v", logMsg)
			o.logMsgChan <- logMsg
		}
		//fmt.Println("MongoOutput:", string(outBytes))
	}

	return
}

func init() {
	pipeline.RegisterPlugin("MongoOutput", func() interface{} {
		return new(MongoOutput)
	})
}

    以上代码已经运行在实际环境中，核心逻辑是将pack Encode后，解析返回数据，转换为Mongo存储的格式，然后将数据传递到mo.logMsgChan，Mongo写routine持续从该channel缓冲中读取数据，将数据写入mongodb中；

3、编译插件到Heka系统，介绍两种方式：

    1）参考官网： http://hekad.readthedocs.org/en/v0.8.2/installing.html#build-include-externals

    2）在调用Heka系统的build.sh后，会产生一个build目录，cd build/heka目录，发现目录结构符合标准的go工程目录，在该目录下建立install文件，install文件如下：

#!/usr/bin/env bash

if [ ! -f install ]; then
echo 'install must be run within its container folder' 1>&2
exit 1
fi

CURDIR=`pwd`
OLDGOPATH="$GOPATH"
export GOPATH="$OLDGOPATH:$CURDIR"

gofmt -w src

#go install hekad
go install github.com/mozilla-services/heka/cmd/hekad
#go install github.com/mozilla-services/heka/cmd/heka-flood
#go install github.com/mozilla-services/heka/cmd/heka-inject
#go install github.com/mozilla-services/heka/cmd/heka-logstreamer
#go install github.com/mozilla-services/heka/cmd/heka-cat
#go install github.com/mozilla-services/heka/cmd/heka-sbmgr
#go install github.com/mozilla-services/heka/cmd/heka-sbmgrload
#go test    github.com/mozilla-services/heka/plugins/mongo

export GOPATH="$OLDGOPATH"

echo 'finished'

    每次在添加完插件后，需要在hekad/main.go文件import中添加一行代码：

_ "github.com/mozilla-services/heka/plugins/mongo"

    ok之后，编译./install即可！

以上如有疑问，请email到cxwshawn@yeah.net

转载于:https://my.oschina.net/shawnChen/blog/361152