registrar.go

package crawler

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	cfg "github.com/elastic/beats/filebeat/config"
	"github.com/elastic/beats/filebeat/input"
	. "github.com/elastic/beats/filebeat/input"
	"github.com/elastic/beats/libbeat/logp"
)

type Registrar struct {
	// Path to the Registry File
	registryFile string
	// Map with all file paths inside and the corresponding state
	state      map[string]*FileState
	stateMutex sync.Mutex
	// Channel used by the prospector and crawler to send FileStates to be persisted
	Persist chan *input.FileState
	running bool

	Channel chan []*FileEvent
	done    chan struct{}
}

func NewRegistrar(registryFile string) (*Registrar, error) {

	r := &Registrar{
		registryFile: registryFile,
		done:         make(chan struct{}),
	}
	err := r.Init()

	return r, err
}

func (r *Registrar) Init() error {
	// Init state
	r.Persist = make(chan *FileState)
	r.state = map[string]*FileState{}
	r.Channel = make(chan []*FileEvent, 1)

	// Set to default in case it is not set
	if r.registryFile == "" {
		r.registryFile = cfg.DefaultRegistryFile
	}

	// Make sure the directory where we store the registryFile exists
	absPath, err := filepath.Abs(r.registryFile)
	if err != nil {
		return fmt.Errorf("Failed to get the absolute path of %s: %v",
			r.registryFile, err)
	}
	r.registryFile = absPath

	// Create directory if it does not already exist.
	registryPath := filepath.Dir(r.registryFile)
	err = os.MkdirAll(registryPath, 0755)
	if err != nil {
		return fmt.Errorf("Failed to created registry file dir %s: %v",
			registryPath, err)
	}

	logp.Info("Registry file set to: %s", r.registryFile)

	return nil
}

// loadState fetches the previous reading state from the configure RegistryFile file
// The default file is .filebeat file which is stored in the same path as the binary is running
func (r *Registrar) LoadState() {
	r.stateMutex.Lock()
	defer r.stateMutex.Unlock()
	if existing, e := os.Open(r.registryFile); e == nil {
		defer existing.Close()
		logp.Info("Loading registrar data from %s", r.registryFile)
		decoder := json.NewDecoder(existing)
		decoder.Decode(&r.state)
	}
}

func (r *Registrar) Run() {
	logp.Info("Starting Registrar")

	r.running = true

	// Writes registry on shutdown
	defer r.writeRegistry()

	for {
		select {
		case <-r.done:
			logp.Info("Ending Registrar")
			return
		// Treats new log files to persist with higher priority then new events
		case state := <-r.Persist:
			r.setState(state)
			logp.Debug("prospector", "Registrar will re-save state for %s", state.Source)
		case events := <-r.Channel:
			r.processEvents(events)
		}

		if e := r.writeRegistry(); e != nil {
			// REVU: but we should panic, or something, right?
			logp.Err("Writing of registry returned error: %v. Continuing..", e)
		}
	}
}

func (r *Registrar) processEvents(events []*FileEvent) {
	logp.Debug("registrar", "Processing %d events", len(events))

	// Take the last event found for each file source
	for _, event := range events {
		if !r.running {
			break
		}

		// skip stdin
		if event.InputType == cfg.StdinInputType {
			continue
		}

		r.setState(event.GetState())
	}
}

func (r *Registrar) Stop() {
	logp.Info("Stopping Registrar")
	r.running = false
	close(r.done)
	// Note: don't block using waitGroup, cause this method is run by async signal handler
}

func (r *Registrar) GetFileState(path string) (*FileState, bool) {
	state, exist := r.getState(path)
	return state, exist
}

// writeRegistry writes the new json registry file to disk.
func (r *Registrar) writeRegistry() error {
	logp.Debug("registrar", "Write registry file: %s", r.registryFile)

	tempfile := r.registryFile + ".new"
	file, e := os.Create(tempfile)
	if e != nil {
		logp.Err("Failed to create tempfile (%s) for writing: %s", tempfile, e)
		return e
	}

	encoder := json.NewEncoder(file)
	state := r.getStateCopy()
	encoder.Encode(state)

	// Directly close file because of windows
	file.Close()

	logp.Info("Registry file updated. %d states written.", len(state))

	return SafeFileRotate(r.registryFile, tempfile)
}

func (r *Registrar) fetchState(filePath string, fileInfo os.FileInfo) (int64, bool) {

	// Check if there is a state for this file
	lastState, isFound := r.GetFileState(filePath)

	// Use os info to compare files as old fileInfo not necessary available
	if isFound && lastState.FileStateOS.IsSame(GetOSFileState(&fileInfo)) {

		logp.Debug("registar", "Same file as before found. Fetch the state.")
		// We're resuming - throw the last state back downstream so we resave it
		// And return the offset - also force harvest in case the file is old and we're about to skip it
		return lastState.Offset, true
	}

	if previous, err := r.getPreviousFile(filePath, fileInfo); err == nil {
		// File has rotated between shutdown and startup
		// We return last state downstream, with a modified event source with the new file name
		// And return the offset - also force harvest in case the file is old and we're about to skip it
		logp.Info("Detected rename of a previously harvested file: %s -> %s", previous, filePath)

		lastState, _ := r.GetFileState(previous)
		return lastState.Offset, true
	}

	if isFound {
		logp.Info("Not resuming rotated file: %s", filePath)
	}

	// New file so just start from an automatic position
	return 0, false
}

// getPreviousFile checks in the registrar if there is the newFile already exist with a different name
// In case an old file is found, the path to the file is returned, if not, an error is returned
func (r *Registrar) getPreviousFile(newFilePath string, newFileInfo os.FileInfo) (string, error) {

	newState := input.GetOSFileState(&newFileInfo)

	for oldFilePath, oldState := range r.getStateCopy() {

		// Skipping when path the same
		if oldFilePath == newFilePath {
			continue
		}

		// Compare states
		if newState.IsSame(oldState.FileStateOS) {
			logp.Info("Old file with new name found: %s is no %s", oldFilePath, newFilePath)
			return oldFilePath, nil
		}
	}

	return "", fmt.Errorf("No previous file found")
}

func (r *Registrar) setState(state *FileState) {
	r.stateMutex.Lock()
	defer r.stateMutex.Unlock()
	r.state[state.Source] = state
}

func (r *Registrar) getState(path string) (*FileState, bool) {
	r.stateMutex.Lock()
	defer r.stateMutex.Unlock()
	state, exist := r.state[path]
	return state, exist
}

func (r *Registrar) getStateCopy() map[string]FileState {
	r.stateMutex.Lock()
	defer r.stateMutex.Unlock()

	copy := make(map[string]FileState)
	for k, v := range r.state {
		copy[k] = *v
	}

	return copy
}