Golang toml.PrimitiveDecode函数代码示例

// OrigPrepCommonTypedConfig is the default implementation of the
// `PrepCommonTypedConfig` function.
func (m *pluginMaker) OrigPrepCommonTypedConfig() (interface{}, error) {
	var (
		commonTypedConfig interface{}
		err               error
	)
	switch m.category {
	case "Input":
		commonInput := CommonInputConfig{
			Retries: getDefaultRetryOptions(),
		}
		err = toml.PrimitiveDecode(m.tomlSection, &commonInput)
		commonTypedConfig = commonInput
	case "Filter", "Output":
		commonFO := CommonFOConfig{
			Retries: getDefaultRetryOptions(),
		}
		err = toml.PrimitiveDecode(m.tomlSection, &commonFO)
		commonTypedConfig = commonFO
	case "Splitter":
		commonSplitter := CommonSplitterConfig{}
		err = toml.PrimitiveDecode(m.tomlSection, &commonSplitter)
		commonTypedConfig = commonSplitter
	}
	if err != nil {
		return nil, err
	}
	return commonTypedConfig, nil
}

// If `configable` supports the `HasConfigStruct` interface this will use said
// interface to fetch a config struct object and populate it w/ the values in
// provided `config`. If not, simply returns `config` unchanged.
func LoadConfigStruct(config toml.Primitive, configable interface{}) (
	configStruct interface{}, err error) {

	// On two lines for scoping reasons.
	hasConfigStruct, ok := configable.(HasConfigStruct)
	if !ok {
		// If we don't have a config struct, change it to a PluginConfig
		configStruct = PluginConfig{}
		if err = toml.PrimitiveDecode(config, configStruct); err != nil {
			configStruct = nil
		}
		return
	}

	defer func() {
		// Slight protection against ConfigStruct call into plugin code.
		if r := recover(); r != nil {
			configStruct = nil
			err = fmt.Errorf("ConfigStruct() panicked: %s", r)
		}
	}()

	configStruct = hasConfigStruct.ConfigStruct()
	if err = toml.PrimitiveDecode(config, configStruct); err != nil {
		configStruct = nil
		err = fmt.Errorf("Can't unmarshal config: %s", err)
	}
	return
}

// NewPluginMaker creates and returns a PluginMaker that can generate running
// plugins for the provided TOML configuration. It will load the plugin type
// and extract any of the Heka-defined common config for the plugin before
// returning.
func NewPluginMaker(name string, pConfig *PipelineConfig, tomlSection toml.Primitive) (
	PluginMaker, error) {

	// Create the maker, extract plugin type, and make sure the plugin type
	// exists.
	maker := &pluginMaker{
		name:         name,
		tomlSection:  tomlSection,
		commonConfig: CommonConfig{},
		pConfig:      pConfig,
	}

	var err error
	if err = toml.PrimitiveDecode(tomlSection, &maker.commonConfig); err != nil {
		return nil, fmt.Errorf("can't decode common config for '%s': %s", name, err)
	}
	if maker.commonConfig.Typ == "" {
		maker.commonConfig.Typ = name
	}
	constructor, ok := AvailablePlugins[maker.commonConfig.Typ]
	if !ok {
		return nil, fmt.Errorf("No registered plugin type: %s", maker.commonConfig.Typ)
	}
	maker.constructor = constructor

	// Extract plugin category and any category-specific common (i.e. Heka
	// defined) configuration.
	maker.category = getPluginCategory(maker.commonConfig.Typ)
	if maker.category == "" {
		return nil, errors.New("Unrecognized plugin category")
	}

	switch maker.category {
	case "Input":
		commonInput := CommonInputConfig{
			Retries: getDefaultRetryOptions(),
		}
		err = toml.PrimitiveDecode(tomlSection, &commonInput)
		maker.commonTypedConfig = commonInput
	case "Filter", "Output":
		commonFO := CommonFOConfig{
			Retries: getDefaultRetryOptions(),
		}
		err = toml.PrimitiveDecode(tomlSection, &commonFO)
		maker.commonTypedConfig = commonFO
	case "Splitter":
		commonSplitter := CommonSplitterConfig{}
		err = toml.PrimitiveDecode(tomlSection, &commonSplitter)
		maker.commonTypedConfig = commonSplitter
	}

	if err != nil {
		return nil, fmt.Errorf("can't decode common %s config for '%s': %s",
			strings.ToLower(maker.category), name, err)
	}
	return maker, nil
}

func (this *iRunner) Start(conf toml.Primitive) {
	plugCommon := &PluginCommonConfig{
		Type: "",
		Tag:  "",
	}
	if err := toml.PrimitiveDecode(conf, plugCommon); err != nil {
		log.Fatalln("toml struct error")
	}

	input, ok := input_plugins[plugCommon.Type]
	if !ok {
		log.Fatalln("unkown type ", plugCommon.Type)
	}

	in := input()

	err := in.(Input).Init(plugCommon, conf)
	if err != nil {
		log.Fatalln("in.(Input).Init", err)
	}

	err = in.(Input).Run(this)
	if err != nil {
		log.Fatalln("in.(Input).Run", err)
	}
}

// Loads the config for a section supplied, configures the supplied object, and initializes
func LoadConfigForSection(app *Application, sectionName string, obj HasConfigStruct,
	env envconf.Environment, configFile ConfigFile) (err error) {

	conf, ok := configFile[sectionName]
	if !ok {
		return fmt.Errorf("Error loading config file, section: %s", sectionName)
	}
	confStruct := obj.ConfigStruct()
	if confStruct == nil {
		return nil
	}

	if err = toml.PrimitiveDecode(conf, confStruct); err != nil {
		return fmt.Errorf("Unable to decode config for section '%s': %s",
			sectionName, err)
	}

	if err = env.Decode(toEnvName(sectionName), EnvSep, confStruct); err != nil {
		return fmt.Errorf("Invalid environment variable for section '%s': %s",
			sectionName, err)
	}

	err = obj.Init(app, confStruct)
	return
}

func (self *TcpInput) Init(pcf *plugins.PluginCommonConfig, conf toml.Primitive) (err error) {

	log.Println("TcpInput Init")
	self.common = pcf
	self.config = &TcpInputConfig{
		Net: "tcp",
	}
	if err := toml.PrimitiveDecode(conf, self.config); err != nil {
		return fmt.Errorf("Can't unmarshal TcpInput config: %s", err)
	}
	address, err := net.ResolveTCPAddr(self.config.Net, self.config.Address)
	if err != nil {
		return fmt.Errorf("ResolveTCPAddress failed: %s\n", err.Error())
	}
	self.listener, err = net.ListenTCP(self.config.Net, address)
	if err != nil {
		return fmt.Errorf("ListenTCP failed: %s\n", err.Error())
	}
	// We're already listening, make sure we clean up if init fails later on.
	closeIt := true
	defer func() {
		if closeIt {
			self.listener.Close()
		}
	}()
	if self.config.KeepAlivePeriod != 0 {
		self.keepAliveDuration = time.Duration(self.config.KeepAlivePeriod) * time.Second
	}
	self.stopChan = make(chan bool)
	closeIt = false
	return nil
}

func main() {
	configFile := flag.String("config", "logstreamer.toml", "Heka Logstreamer configuration file")

	flag.Parse()

	if flag.NFlag() == 0 {
		flag.PrintDefaults()
		os.Exit(0)
	}

	fconfig := make(FileConfig)
	if _, err := toml.DecodeFile(*configFile, &fconfig); err != nil {
		log.Printf("Error decoding config file: %s", err)
		return
	}

	// Filter out logstream inputs
	inputs := make(map[string]toml.Primitive)
	for name, prim := range fconfig {
		basic := new(Basic)
		if name == "LogstreamerInput" {
			inputs[name] = prim
		} else if err := toml.PrimitiveDecode(prim, &basic); err == nil {
			if basic.PluginType == "LogstreamerInput" {
				inputs[name] = prim
			}
		}
	}

	// Go through the logstreams and parse their configs
	for name, prim := range inputs {
		parseConfig(name, prim)
	}
}

func (this *TailsInput) Init(pcf *plugins.PluginCommonConfig, conf toml.Primitive) (err error) {
	this.common = pcf
	this.config = &TailsInputConfig{
		LogDirectory:     "/var/log",
		JournalDirectory: "/tmp/",
		//FileMatch: "*.log",
		RescanInterval: "1m",
		SyncInterval:   2,
	}
	this.files = make([]string, 0)
	if err := toml.PrimitiveDecode(conf, this.config); err != nil {
		return fmt.Errorf("Can't unmarshal tails config: %s", err)
	}
	if this.config.FileMatch == "" {
		return errors.New("`file_match` setting is required.")
	}
	if len(this.config.FileMatch) > 0 && this.config.FileMatch[len(this.config.FileMatch)-1:] != "$" {
		this.config.FileMatch += "$"
	}
	// Setup the rescan interval.
	if this.rescanInterval, err = time.ParseDuration(this.config.RescanInterval); err != nil {
		return
	}
	if !fileExists(this.config.JournalDirectory) {
		if err = os.MkdirAll(filepath.Dir(this.config.JournalDirectory), 0766); err != nil {
			return
		}
	}
	return nil
}

// Creates a FilterRunner for the specified sandbox name and configuration
func createRunner(dir, name string, configSection toml.Primitive) (FilterRunner, error) {
	var err error
	var pluginGlobals PluginGlobals

	wrapper := new(PluginWrapper)
	wrapper.name = name

	pluginGlobals.Retries = RetryOptions{
		MaxDelay:   "30s",
		Delay:      "250ms",
		MaxRetries: -1,
	}

	if err = toml.PrimitiveDecode(configSection, &pluginGlobals); err != nil {
		return nil, fmt.Errorf("Unable to decode config for plugin: %s, error: %s",
			wrapper.name, err.Error())
	}
	if pluginGlobals.Typ != "SandboxFilter" {
		return nil, fmt.Errorf("Plugin must be a SandboxFilter, received %s",
			pluginGlobals.Typ)
	}

	// Create plugin, test config object generation.
	wrapper.pluginCreator, _ = AvailablePlugins[pluginGlobals.Typ]
	plugin := wrapper.pluginCreator()
	var config interface{}
	if config, err = LoadConfigStruct(configSection, plugin); err != nil {
		return nil, fmt.Errorf("Can't load config for '%s': %s", wrapper.name, err)
	}
	wrapper.configCreator = func() interface{} { return config }
	conf := config.(*sandbox.SandboxConfig)
	conf.ScriptFilename = filepath.Join(dir, fmt.Sprintf("%s.%s", wrapper.name, conf.ScriptType))
	if wantsName, ok := plugin.(WantsName); ok {
		wantsName.SetName(wrapper.name)
	}

	// Apply configuration to instantiated plugin.
	if err = plugin.(Plugin).Init(config); err != nil {
		return nil, fmt.Errorf("Initialization failed for '%s': %s", name, err)
	}

	runner := NewFORunner(wrapper.name, plugin.(Plugin), &pluginGlobals)
	runner.name = wrapper.name

	if pluginGlobals.Ticker != 0 {
		runner.tickLength = time.Duration(pluginGlobals.Ticker) * time.Second
	}

	var matcher *MatchRunner
	if pluginGlobals.Matcher != "" {
		if matcher, err = NewMatchRunner(pluginGlobals.Matcher,
			pluginGlobals.Signer, runner); err != nil {
			return nil, fmt.Errorf("Can't create message matcher for '%s': %s",
				wrapper.name, err)
		}
		runner.matcher = matcher
	}

	return runner, nil
}

func (this *Pipeline) LoadConfig(plugConfig map[string]toml.Primitive) error {
	for k, v := range plugConfig {
		log.Printf("v %+v", v)
		plugCommon := &PluginCommonConfig{}
		if err := toml.PrimitiveDecode(v, plugCommon); err != nil {
			return fmt.Errorf("Can't unmarshal config: %s", err)
		}
		pluginType := getPluginType(plugCommon.Type)
		if pluginType == "" {
			continue
		}
		if plugCommon.Tag == "" {
			log.Println("Tag empty")
		}
		switch pluginType {
		case "Input":
			this.InputRunners = append(this.InputRunners, v)
		case "Output":
			this.OutputRunners = append(this.OutputRunners, v)
		case "Encoder":
			this.EncodeRunners = append(this.EncodeRunners, v)
		case "Decoder":
			this.DecodeRunners = append(this.DecodeRunners, v)
		}
		log.Printf("%s => %s", k, plugCommon.Type)
	}

	return nil
}

// Load an extensible section that has a type keyword
func LoadExtensibleSection(app *Application, sectionName string,
	extensions AvailableExtensions, env envconf.Environment,
	configFile ConfigFile) (obj HasConfigStruct, err error) {

	confSection := new(ExtensibleGlobals)

	conf, ok := configFile[sectionName]
	if !ok {
		return nil, fmt.Errorf("Missing section '%s'", sectionName)
	}

	if err = toml.PrimitiveDecode(conf, confSection); err != nil {
		return nil, err
	}
	if err = env.Decode(toEnvName(sectionName), EnvSep, confSection); err != nil {
		return nil, err
	}
	ext, ok := extensions.Get(confSection.Typ)
	if !ok {
		//TODO: Add log info to indicate using "default"
		return nil, fmt.Errorf("No type '%s' available to load for section '%s'",
			confSection.Typ, sectionName)
	}

	obj = ext()
	loadedConfig, err := LoadConfigStruct(sectionName, env, conf, obj)
	if err != nil {
		return nil, err
	}

	err = obj.Init(app, loadedConfig)
	return obj, err
}

func (self *UdpInput) Init(pcf *plugins.PluginCommonConfig, conf toml.Primitive) (err error) {

	log.Println("UdpInput Init")
	self.common = pcf
	self.config = &UdpInputConfig{
		Net: "udp4",
	}
	if err := toml.PrimitiveDecode(conf, self.config); err != nil {
		return fmt.Errorf("Can't unmarshal UdpInput config: %s", err)
	}
	address, err := net.ResolveUDPAddr(self.config.Net, self.config.Address)
	if err != nil {
		return fmt.Errorf("ResolveUDPAddress failed: %s\n", err.Error())
	}
	self.listener, err = net.ListenUDP(self.config.Net, address)
	if err != nil {
		return fmt.Errorf("ListenUDP failed: %s\n", err.Error())
	}
	self.listener.SetReadBuffer(1048576)
	closeIt := true
	defer func() {
		if closeIt {
			self.listener.Close()
		}
	}()
	self.stopChan = make(chan bool)
	self.inChan = make(chan UdpPack, 1024)
	closeIt = false
	return nil
}

// loadSection must be passed a plugin name and the config for that plugin. It
// will create a PluginWrapper (i.e. a factory).
func (md *MultiDecoder) loadSection(sectionName string,
	configSection toml.Primitive) (plugin Decoder, err error) {
	var ok bool
	var pluginGlobals PluginGlobals
	var pluginType string

	wrapper := new(PluginWrapper)
	wrapper.Name = sectionName

	// Setup default retry policy
	pluginGlobals.Retries = RetryOptions{
		MaxDelay:   "30s",
		Delay:      "250ms",
		MaxRetries: -1,
	}

	if err = toml.PrimitiveDecode(configSection, &pluginGlobals); err != nil {
		err = fmt.Errorf("%s Unable to decode config for plugin: %s, error: %s", md.Name, wrapper.Name, err.Error())
		md.log(err.Error())
		return
	}

	if pluginGlobals.Typ == "" {
		pluginType = sectionName
	} else {
		pluginType = pluginGlobals.Typ
	}

	if wrapper.PluginCreator, ok = AvailablePlugins[pluginType]; !ok {
		err = fmt.Errorf("%s No such plugin: %s (type: %s)", md.Name, wrapper.Name, pluginType)
		md.log(err.Error())
		return
	}

	// Create plugin, test config object generation.
	plugin = wrapper.PluginCreator().(Decoder)
	var config interface{}
	if config, err = LoadConfigStruct(configSection, plugin); err != nil {
		err = fmt.Errorf("%s Can't load config for %s '%s': %s", md.Name,
			sectionName,
			wrapper.Name, err)
		md.log(err.Error())
		return
	}
	wrapper.ConfigCreator = func() interface{} { return config }

	if wantsName, ok := plugin.(WantsName); ok {
		wantsName.SetName(wrapper.Name)
	}

	// Apply configuration to instantiated plugin.
	if err = plugin.(Plugin).Init(config); err != nil {
		err = fmt.Errorf("Initialization failed for '%s': %s",
			sectionName, err)
		md.log(err.Error())
		return
	}
	return
}

// PrepConfig generates a config struct for the plugin (instantiating an
// instance of the plugin to do so, if necessary) and decodes the TOML config
// into the generated struct.
func (m *pluginMaker) PrepConfig() error {
	if m.configPrepped {
		// Already done, just return.
		return nil
	}

	if m.configStruct == nil {
		m.makeConfig()
	} else if m.plugin == nil {
		m.plugin = m.makePlugin()
	}

	if _, ok := m.plugin.(HasConfigStruct); !ok {
		// If plugin doesn't implement HasConfigStruct then we're decoding
		// into an empty PluginConfig object.
		if err := toml.PrimitiveDecode(m.tomlSection, m.configStruct); err != nil {
			return fmt.Errorf("can't decode config for '%s': %s ", m.name, err.Error())
		}
		m.configPrepped = true
		return nil
	}

	// Use reflection to extract the fields (or TOML tag names, if available)
	// of the values that Heka has already extracted so we know they're not
	// required to be specified in the config struct.
	hekaParams := make(map[string]interface{})
	commons := []interface{}{m.commonConfig, m.commonTypedConfig}
	for _, common := range commons {
		if common == nil {
			continue
		}
		rt := reflect.ValueOf(common).Type()
		for i := 0; i < rt.NumField(); i++ {
			sft := rt.Field(i)
			kname := sft.Tag.Get("toml")
			if len(kname) == 0 {
				kname = sft.Name
			}
			hekaParams[kname] = true
		}
	}

	// Finally decode the TOML into the struct. Use of PrimitiveDecodeStrict
	// means that an error will be raised for any config options in the TOML
	// that don't have corresponding attributes on the struct, delta the
	// hekaParams that can be safely excluded.
	err := toml.PrimitiveDecodeStrict(m.tomlSection, m.configStruct, hekaParams)
	if err != nil {
		matches := unknownOptionRegex.FindStringSubmatch(err.Error())
		if len(matches) == 2 {
			// We've got an unrecognized config option.
			return fmt.Errorf("unknown config setting for '%s': %s", m.name, matches[1])
		}
		return err
	}

	m.configPrepped = true
	return nil
}

func main() {
	configFile := flag.String("config", "logstreamer.toml", "Heka Logstreamer configuration file")

	flag.Parse()

	if flag.NFlag() == 0 {
		flag.PrintDefaults()
		os.Exit(0)
	}

	p, err := os.Open(*configFile)
	if err != nil {
		client.LogError.Fatalf("Error opening config file: %s", err)
	}
	fi, err := p.Stat()
	if err != nil {
		client.LogError.Fatalf("Error fetching config file info: %s", err)
	}

	fconfig := make(FileConfig)
	if fi.IsDir() {
		files, _ := ioutil.ReadDir(*configFile)
		for _, f := range files {
			fName := f.Name()
			if strings.HasPrefix(fName, ".") || strings.HasSuffix(fName, ".bak") ||
				strings.HasSuffix(fName, ".tmp") || strings.HasSuffix(fName, "~") {
				// Skip obviously non-relevant files.
				continue
			}
			fPath := filepath.Join(*configFile, fName)
			if _, err = toml.DecodeFile(fPath, &fconfig); err != nil {
				client.LogError.Fatalf("Error decoding config file: %s", err)
			}
		}
	} else {
		if _, err := toml.DecodeFile(*configFile, &fconfig); err != nil {
			client.LogError.Fatalf("Error decoding config file: %s", err)
		}
	}

	// Filter out logstream inputs
	inputs := make(map[string]toml.Primitive)
	for name, prim := range fconfig {
		basic := new(Basic)
		if name == "LogstreamerInput" {
			inputs[name] = prim
		} else if err := toml.PrimitiveDecode(prim, &basic); err == nil {
			if basic.PluginType == "LogstreamerInput" {
				inputs[name] = prim
			}
		}
	}

	// Go through the logstreams and parse their configs
	for name, prim := range inputs {
		parseConfig(name, prim)
	}
}