285 lines
7.7 KiB
Go
285 lines
7.7 KiB
Go
|
// Licensed to Elasticsearch B.V. under one or more contributor
|
||
|
|
// license agreements. See the NOTICE file distributed with
|
||
|
|
// this work for additional information regarding copyright
|
||
|
|
// ownership. Elasticsearch B.V. licenses this file to you under
|
||
|
|
// the Apache License, Version 2.0 (the "License"); you may
|
||
|
|
// not use this file except in compliance with the License.
|
||
|
|
// You may obtain a copy of the License at
|
||
|
|
//
|
||
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
//
|
||
|
|
// Unless required by applicable law or agreed to in writing,
|
||
|
|
// software distributed under the License is distributed on an
|
||
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||
|
|
// KIND, either express or implied. See the License for the
|
||
|
|
// specific language governing permissions and limitations
|
||
|
|
// under the License.
|
||
|
|
|
||
|
|
// Package checkpoint persists event log state information to disk so that
|
||
|
|
// event log monitoring can resume from the last read event in the case of a
|
||
|
|
// restart or unexpected interruption.
|
||
|
|
package checkpoint
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"io/ioutil"
|
||
|
|
"os"
|
||
|
|
"path/filepath"
|
||
|
|
"sort"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"gopkg.in/yaml.v2"
|
||
|
|
|
||
|
|
"github.com/elastic/beats/libbeat/logp"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Checkpoint persists event log state information to disk.
|
||
|
|
type Checkpoint struct {
|
||
|
|
wg sync.WaitGroup // WaitGroup used to wait on the shutdown of the checkpoint worker.
|
||
|
|
done chan struct{} // Channel for shutting down the checkpoint worker.
|
||
|
|
once sync.Once // Used to guarantee shutdown happens once.
|
||
|
|
file string // File where the state is persisted.
|
||
|
|
numUpdates int // Number of updates received since last persisting to disk.
|
||
|
|
maxUpdates int // Maximum number of updates to buffer before persisting to disk.
|
||
|
|
flushInterval time.Duration // Maximum time interval that can pass before persisting to disk.
|
||
|
|
sort []string // Slice used for sorting states map (store to save on mallocs).
|
||
|
|
|
||
|
|
lock sync.RWMutex
|
||
|
|
states map[string]EventLogState
|
||
|
|
|
||
|
|
save chan EventLogState
|
||
|
|
}
|
||
|
|
|
||
|
|
// PersistedState represents the format of the data persisted to disk.
|
||
|
|
type PersistedState struct {
|
||
|
|
UpdateTime time.Time `yaml:"update_time"`
|
||
|
|
States []EventLogState `yaml:"event_logs"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// EventLogState represents the state of an individual event log.
|
||
|
|
type EventLogState struct {
|
||
|
|
Name string `yaml:"name"`
|
||
|
|
RecordNumber uint64 `yaml:"record_number"`
|
||
|
|
Timestamp time.Time `yaml:"timestamp"`
|
||
|
|
Bookmark string `yaml:"bookmark,omitempty"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewCheckpoint creates and returns a new Checkpoint. This method loads state
|
||
|
|
// information from disk if it exists and starts a goroutine for persisting
|
||
|
|
// state information to disk. Shutdown should be called when finished to
|
||
|
|
// guarantee any in-memory state information is flushed to disk.
|
||
|
|
//
|
||
|
|
// file is the name of the file where event log state is persisted as YAML.
|
||
|
|
// maxUpdates is the maximum number of updates checkpoint will accept before
|
||
|
|
// triggering a flush to disk. interval is maximum amount of time that can
|
||
|
|
// pass since the last flush before triggering a flush to disk (minimum value
|
||
|
|
// is 1s).
|
||
|
|
func NewCheckpoint(file string, maxUpdates int, interval time.Duration) (*Checkpoint, error) {
|
||
|
|
c := &Checkpoint{
|
||
|
|
done: make(chan struct{}),
|
||
|
|
file: file,
|
||
|
|
maxUpdates: maxUpdates,
|
||
|
|
flushInterval: interval,
|
||
|
|
sort: make([]string, 0, 10),
|
||
|
|
states: make(map[string]EventLogState),
|
||
|
|
save: make(chan EventLogState, 1),
|
||
|
|
}
|
||
|
|
|
||
|
|
// Minimum batch size.
|
||
|
|
if c.maxUpdates < 1 {
|
||
|
|
c.maxUpdates = 1
|
||
|
|
}
|
||
|
|
|
||
|
|
// Minimum flush interval.
|
||
|
|
if c.flushInterval < time.Second {
|
||
|
|
c.flushInterval = time.Second
|
||
|
|
}
|
||
|
|
|
||
|
|
// Read existing state information:
|
||
|
|
ps, err := c.read()
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
if ps != nil {
|
||
|
|
for _, state := range ps.States {
|
||
|
|
c.states[state.Name] = state
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Write the state file to verify we have have permissions.
|
||
|
|
err = c.flush()
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
c.wg.Add(1)
|
||
|
|
go c.run()
|
||
|
|
return c, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// run is worker loop that reads incoming state information from the save
|
||
|
|
// channel and persists it when the number of changes reaches maxEvents or
|
||
|
|
// the amount of time since the last disk write reaches flushInterval.
|
||
|
|
func (c *Checkpoint) run() {
|
||
|
|
defer c.wg.Done()
|
||
|
|
defer c.persist()
|
||
|
|
|
||
|
|
flushTimer := time.NewTimer(c.flushInterval)
|
||
|
|
defer flushTimer.Stop()
|
||
|
|
loop:
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case <-c.done:
|
||
|
|
break loop
|
||
|
|
case s := <-c.save:
|
||
|
|
c.lock.Lock()
|
||
|
|
c.states[s.Name] = s
|
||
|
|
c.lock.Unlock()
|
||
|
|
c.numUpdates++
|
||
|
|
if c.numUpdates < c.maxUpdates {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
case <-flushTimer.C:
|
||
|
|
}
|
||
|
|
|
||
|
|
c.persist()
|
||
|
|
flushTimer.Reset(c.flushInterval)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Shutdown stops the checkpoint worker (which persists any state to disk as
|
||
|
|
// it stops). This method blocks until the checkpoint worker shutdowns. Calling
|
||
|
|
// this method more once is safe and has no effect.
|
||
|
|
func (c *Checkpoint) Shutdown() {
|
||
|
|
c.once.Do(func() {
|
||
|
|
close(c.done)
|
||
|
|
c.wg.Wait()
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
// States returns the current in-memory event log state. This state information
|
||
|
|
// is bootstrapped with any data found on disk at creation time.
|
||
|
|
func (c *Checkpoint) States() map[string]EventLogState {
|
||
|
|
c.lock.RLock()
|
||
|
|
defer c.lock.RUnlock()
|
||
|
|
|
||
|
|
copy := make(map[string]EventLogState)
|
||
|
|
for k, v := range c.states {
|
||
|
|
copy[k] = v
|
||
|
|
}
|
||
|
|
|
||
|
|
return copy
|
||
|
|
}
|
||
|
|
|
||
|
|
// Persist queues the given event log state information to be written to disk.
|
||
|
|
func (c *Checkpoint) Persist(name string, recordNumber uint64, ts time.Time, bookmark string) {
|
||
|
|
c.PersistState(EventLogState{
|
||
|
|
Name: name,
|
||
|
|
RecordNumber: recordNumber,
|
||
|
|
Timestamp: ts,
|
||
|
|
Bookmark: bookmark,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
// PersistState queues the given event log state to be written to disk.
|
||
|
|
func (c *Checkpoint) PersistState(st EventLogState) {
|
||
|
|
c.save <- st
|
||
|
|
}
|
||
|
|
|
||
|
|
// persist writes the current state to disk if the in-memory state is dirty.
|
||
|
|
func (c *Checkpoint) persist() bool {
|
||
|
|
if c.numUpdates == 0 {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
err := c.flush()
|
||
|
|
if err != nil {
|
||
|
|
logp.Err("%v", err)
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
logp.Debug("checkpoint", "Checkpoint saved to disk. numUpdates=%d",
|
||
|
|
c.numUpdates)
|
||
|
|
c.numUpdates = 0
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
// flush writes the current state to disk.
|
||
|
|
func (c *Checkpoint) flush() error {
|
||
|
|
tempFile := c.file + ".new"
|
||
|
|
file, err := create(tempFile)
|
||
|
|
if os.IsNotExist(err) {
|
||
|
|
// Try to create directory if it does not exist.
|
||
|
|
if createDirErr := c.createDir(); createDirErr == nil {
|
||
|
|
file, err = create(tempFile)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("Failed to flush state to disk. %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Sort persisted eventLogs by name.
|
||
|
|
c.sort = c.sort[:0]
|
||
|
|
for k := range c.states {
|
||
|
|
c.sort = append(c.sort, k)
|
||
|
|
}
|
||
|
|
sort.Strings(c.sort)
|
||
|
|
|
||
|
|
ps := PersistedState{
|
||
|
|
UpdateTime: time.Now().UTC(),
|
||
|
|
States: make([]EventLogState, len(c.sort)),
|
||
|
|
}
|
||
|
|
for i, name := range c.sort {
|
||
|
|
ps.States[i] = c.states[name]
|
||
|
|
}
|
||
|
|
|
||
|
|
data, err := yaml.Marshal(ps)
|
||
|
|
if err != nil {
|
||
|
|
file.Close()
|
||
|
|
return fmt.Errorf("Failed to flush state to disk. Could not marshal "+
|
||
|
|
"data to YAML. %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
_, err = file.Write(data)
|
||
|
|
if err != nil {
|
||
|
|
file.Close()
|
||
|
|
return fmt.Errorf("Failed to flush state to disk. Could not write to "+
|
||
|
|
"%s. %v", tempFile, err)
|
||
|
|
}
|
||
|
|
|
||
|
|
file.Close()
|
||
|
|
err = os.Rename(tempFile, c.file)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
// read loads the persisted state from disk. If the file does not exists then
|
||
|
|
// the method returns nil and no error.
|
||
|
|
func (c *Checkpoint) read() (*PersistedState, error) {
|
||
|
|
contents, err := ioutil.ReadFile(c.file)
|
||
|
|
if err != nil {
|
||
|
|
if os.IsNotExist(err) {
|
||
|
|
err = nil
|
||
|
|
}
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
ps := &PersistedState{}
|
||
|
|
err = yaml.Unmarshal(contents, ps)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
return ps, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// createDir creates the directory in which the state file will reside if the
|
||
|
|
// directory does not already exist.
|
||
|
|
func (c *Checkpoint) createDir() error {
|
||
|
|
dir := filepath.Dir(c.file)
|
||
|
|
logp.Info("Creating %s if it does not exist.", dir)
|
||
|
|
return os.MkdirAll(dir, os.FileMode(0750))
|
||
|
|
}
|