forked from influxdata/telegraf
/
service.go
145 lines (120 loc) · 3.21 KB
/
service.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
package snapshotter
import (
"encoding"
"encoding/json"
"fmt"
"log"
"net"
"os"
"strings"
"sync"
"github.com/influxdb/influxdb/snapshot"
"github.com/influxdb/influxdb/tsdb"
)
// MuxHeader is the header byte used for the TCP muxer.
const MuxHeader = 3
// Service manages the listener for the snapshot endpoint.
type Service struct {
wg sync.WaitGroup
err chan error
MetaStore interface {
encoding.BinaryMarshaler
}
TSDBStore *tsdb.Store
Listener net.Listener
Logger *log.Logger
}
// NewService returns a new instance of Service.
func NewService() *Service {
return &Service{
err: make(chan error),
Logger: log.New(os.Stderr, "[snapshot] ", log.LstdFlags),
}
}
// Open starts the service.
func (s *Service) Open() error {
s.Logger.Println("Starting snapshot service")
s.wg.Add(1)
go s.serve()
return nil
}
// Close implements the Service interface.
func (s *Service) Close() error {
if s.Listener != nil {
s.Listener.Close()
}
s.wg.Wait()
return nil
}
// SetLogger sets the internal logger to the logger passed in.
func (s *Service) SetLogger(l *log.Logger) {
s.Logger = l
}
// Err returns a channel for fatal out-of-band errors.
func (s *Service) Err() <-chan error { return s.err }
// serve serves snapshot requests from the listener.
func (s *Service) serve() {
defer s.wg.Done()
for {
// Wait for next connection.
conn, err := s.Listener.Accept()
if err != nil && strings.Contains(err.Error(), "connection closed") {
s.Logger.Println("snapshot listener closed")
return
} else if err != nil {
s.Logger.Println("error accepting snapshot request: ", err.Error())
continue
}
// Handle connection in separate goroutine.
s.wg.Add(1)
go func(conn net.Conn) {
defer s.wg.Done()
defer conn.Close()
if err := s.handleConn(conn); err != nil {
s.Logger.Println(err)
}
}(conn)
}
}
// handleConn processes conn. This is run in a separate goroutine.
func (s *Service) handleConn(conn net.Conn) error {
// Read manifest from connection.
m, err := s.readManifest(conn)
if err != nil {
return fmt.Errorf("read manifest: %s", err)
}
// Write snapshot to connection.
if err := s.writeSnapshot(conn, m); err != nil {
return fmt.Errorf("write snapshot: %s", err)
}
return nil
}
// readManifest reads the manifest size and contents from conn.
// Unmarshals the bytes and returns a manifest object.
func (s *Service) readManifest(conn net.Conn) (snapshot.Manifest, error) {
var m snapshot.Manifest
if err := json.NewDecoder(conn).Decode(&m); err != nil {
return m, err
}
return m, nil
}
// writeSnapshot creates a snapshot writer, trims the manifest, and writes to conn.
func (s *Service) writeSnapshot(conn net.Conn, prev snapshot.Manifest) error {
// Retrieve and serialize the current meta data.
buf, err := s.MetaStore.MarshalBinary()
if err != nil {
return fmt.Errorf("marshal meta: %s", err)
}
// Build a snapshot writer.
sw, err := tsdb.NewSnapshotWriter(buf, s.TSDBStore)
if err != nil {
return fmt.Errorf("create snapshot writer: %s", err)
}
// Trim old files from snapshot.
sw.Manifest = sw.Manifest.Diff(&prev)
// Write snapshot out to connection.
if _, err := sw.WriteTo(conn); err != nil {
return fmt.Errorf("write to: %s", err)
}
return nil
}