common: introduce ServiceState

Author: lebauce

common/service.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"net"
 	"strings"
+	"sync/atomic"
 )
 
 // ServiceType describes the service type (analyzer or agent)
@@ -35,9 +36,42 @@ const (
 	AgentService ServiceType = "agent"
 )
 
+// ServiceState describes the state of a service.
+type ServiceState int64
+
+// MarshalJSON marshal the connection state to JSON
+func (s *ServiceState) MarshalJSON() ([]byte, error) {
+	switch *s {
+	case StartingState:
+		return []byte("\"starting\""), nil
+	case RunningState:
+		return []byte("\"running\""), nil
+	case StoppingState:
+		return []byte("\"stopping\""), nil
+	case StoppedState:
+		return []byte("\"stopped\""), nil
+	}
+	return nil, fmt.Errorf("Invalid state: %d", s)
+}
+
+// Store atomatically stores the state
+func (s *ServiceState) Store(state ServiceState) {
+	atomic.StoreInt64((*int64)(s), int64(state))
+}
+
+// Load atomatically loads and returns the state
+func (s *ServiceState) Load() ServiceState {
+	return ServiceState(atomic.LoadInt64((*int64)(s)))
+}
+
+// CompareAndSwap executes the compare-and-swap operation for a state
+func (s *ServiceState) CompareAndSwap(old, new ServiceState) bool {
+	return atomic.CompareAndSwapInt64((*int64)(s), int64(old), int64(new))
+}
+
 const (
 	// StoppedState service stopped
-	StoppedState = iota + 1
+	StoppedState ServiceState = iota + 1
 	// StartingState service starting
 	StartingState
 	// RunningState service running

etcd/election.go

@@ -19,7 +19,6 @@ package etcd
 
 import (
 	"sync"
-	"sync/atomic"
 	"time"
 
 	etcd "github.com/coreos/etcd/client"
@@ -42,7 +41,7 @@ type MasterElector struct {
 	listeners  []common.MasterElectionListener
 	cancel     context.CancelFunc
 	master     bool
-	state      int64
+	state      common.ServiceState
 	wg         sync.WaitGroup
 }
 
@@ -131,8 +130,8 @@ func (le *MasterElector) start(first chan struct{}) {
 	le.wg.Add(1)
 	defer le.wg.Done()
 
-	atomic.StoreInt64(&le.state, common.RunningState)
-	for atomic.LoadInt64(&le.state) == common.RunningState {
+	le.state.Store(common.RunningState)
+	for le.state.Load() == common.RunningState {
 		resp, err := watcher.Next(ctx)
 		if err != nil {
 			logging.GetLogger().Errorf("Error while watching etcd: %s", err.Error())
@@ -191,7 +190,7 @@ func (le *MasterElector) StartAndWait() {
 
 // Stop the election mechanism
 func (le *MasterElector) Stop() {
-	if atomic.CompareAndSwapInt64(&le.state, common.RunningState, common.StoppingState) {
+	if le.state.CompareAndSwap(common.RunningState, common.StoppingState) {
 		le.cancel()
 		le.wg.Wait()
 	}

flow/pcap.go

@@ -20,7 +20,6 @@ package flow
 import (
 	"io"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/google/gopacket"
@@ -39,7 +38,7 @@ type PcapWriter struct {
 // PcapTableFeeder replaies a pcap file
 type PcapTableFeeder struct {
 	sync.WaitGroup
-	state       int64
+	state       common.ServiceState
 	replay      bool
 	r           io.ReadCloser
 	handleRead  *pcapgo.Reader
@@ -49,19 +48,19 @@ type PcapTableFeeder struct {
 
 // Start a pcap injector
 func (p *PcapTableFeeder) Start() {
-	if atomic.CompareAndSwapInt64(&p.state, common.StoppedState, common.RunningState) {
+	if p.state.CompareAndSwap(common.StoppedState, common.RunningState) {
 		p.Add(1)
 		go p.feedFlowTable()
 	}
 }
 
 // Stop a pcap injector
 func (p *PcapTableFeeder) Stop() {
-	if atomic.CompareAndSwapInt64(&p.state, common.RunningState, common.StoppingState) {
-		atomic.StoreInt64(&p.state, common.StoppingState)
+	if p.state.CompareAndSwap(common.RunningState, common.StoppingState) {
+		p.state.Store(common.StoppingState)
 		p.r.Close()
 		p.Wait()
-		atomic.StoreInt64(&p.state, common.StoppedState)
+		p.state.Store(common.StoppedState)
 	}
 }
 
@@ -81,12 +80,12 @@ func (p *PcapTableFeeder) feedFlowTable() {
 		logging.GetLogger().Error(err.Error())
 	}
 
-	atomic.StoreInt64(&p.state, common.RunningState)
-	for atomic.LoadInt64(&p.state) == common.RunningState {
+	p.state.Store(common.RunningState)
+	for p.state.Load() == common.RunningState {
 		logging.GetLogger().Debugf("Reading one pcap packet")
 		data, ci, err := p.handleRead.ReadPacketData()
 		if err != nil {
-			if atomic.LoadInt64(&p.state) == common.RunningState && err != io.EOF {
+			if p.state.Load() == common.RunningState && err != io.EOF {
 				logging.GetLogger().Warningf("Failed to read packet: %s\n", err)
 			}
 			p.r.Close()

flow/probes/gopacket.go

@@ -23,7 +23,6 @@ import (
 	"fmt"
 	"io"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/google/gopacket"
@@ -59,7 +58,7 @@ type GoPacketProbe struct {
 	Ctx         Context
 	n           *graph.Node
 	packetProbe PacketProbe
-	state       int64
+	state       common.ServiceState
 	ifName      string
 	bpfFilter   string
 	nsPath      string
@@ -88,7 +87,7 @@ func (p *GoPacketProbe) updateStats(g *graph.Graph, n *graph.Node, captureStats
 		case <-ticker.C:
 			if stats, err := p.packetProbe.Stats(); err != nil {
 				p.Ctx.Logger.Error(err)
-			} else if atomic.LoadInt64(&p.state) == common.RunningState {
+			} else if p.state.Load() == common.RunningState {
 				g.Lock()
 				g.UpdateMetadata(n, "Captures", func(obj interface{}) bool {
 					captureStats.PacketsDropped = stats.PacketsDropped
@@ -108,7 +107,7 @@ func (p *GoPacketProbe) listen(packetCallback func(gopacket.Packet)) error {
 	packetSource := p.packetProbe.PacketSource()
 
 	var errs int
-	for atomic.LoadInt64(&p.state) == common.RunningState {
+	for p.state.Load() == common.RunningState {
 		packet, err := packetSource.NextPacket()
 		switch err {
 		case nil:
@@ -139,7 +138,7 @@ func (p *GoPacketProbe) listen(packetCallback func(gopacket.Packet)) error {
 // Run starts capturing packet, calling the passed callback for every packet
 // and notifying the flow probe handler when the capture has started
 func (p *GoPacketProbe) Run(packetCallback func(gopacket.Packet), e ProbeEventHandler) error {
-	atomic.StoreInt64(&p.state, common.RunningState)
+	p.state.Store(common.RunningState)
 
 	var nsContext *common.NetNSContext
 	var err error
@@ -199,14 +198,14 @@ func (p *GoPacketProbe) Run(packetCallback func(gopacket.Packet), e ProbeEventHa
 	statsTicker.Stop()
 
 	p.packetProbe.Close()
-	atomic.StoreInt64(&p.state, common.StoppedState)
+	p.state.Store(common.StoppedState)
 
 	return err
 }
 
 // Stop capturing packets
 func (p *GoPacketProbe) Stop() {
-	atomic.StoreInt64(&p.state, common.StoppingState)
+	p.state.Store(common.StoppingState)
 }
 
 // NewGoPacketProbe returns a new Gopacket flow probe. It can use either `pcap` or `afpacket`

flow/probes/pcapsocket.go

@@ -21,7 +21,6 @@ import (
 	"fmt"
 	"net"
 	"sync"
-	"sync/atomic"
 
 	"github.com/skydive-project/skydive/api/types"
 	"github.com/skydive-project/skydive/common"
@@ -33,7 +32,7 @@ import (
 // PcapSocketProbe describes a TCP packet listener that inject packets in a flowtable
 type PcapSocketProbe struct {
 	Ctx       Context
-	state     int64
+	state     common.ServiceState
 	flowTable *flow.Table
 	listener  *net.TCPListener
 	port      int
@@ -49,15 +48,15 @@ type PcapSocketProbeHandler struct {
 }
 
 func (p *PcapSocketProbe) run() {
-	atomic.StoreInt64(&p.state, common.RunningState)
+	p.state.Store(common.RunningState)
 
 	packetSeqChan, _, _ := p.flowTable.Start()
 	defer p.flowTable.Stop()
 
-	for atomic.LoadInt64(&p.state) == common.RunningState {
+	for p.state.Load() == common.RunningState {
 		conn, err := p.listener.Accept()
 		if err != nil {
-			if atomic.LoadInt64(&p.state) == common.RunningState {
+			if p.state.Load() == common.RunningState {
 				p.Ctx.Logger.Errorf("Error while accepting connection: %s", err)
 			}
 			break
@@ -132,7 +131,7 @@ func (p *PcapSocketProbeHandler) UnregisterProbe(n *graph.Node, e ProbeEventHand
 
 	p.Ctx.FTA.Release(probe.flowTable)
 
-	atomic.StoreInt64(&probe.state, common.StoppingState)
+	probe.state.Store(common.StoppingState)
 	probe.listener.Close()
 
 	p.portAllocator.Release(probe.port)

flow/server/flow_server.go

@@ -24,7 +24,6 @@ import (
 	"strconv"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/skydive-project/skydive/common"
@@ -84,7 +83,7 @@ type FlowServerWebSocketConn struct {
 type FlowServer struct {
 	storage            storage.Storage
 	conn               FlowServerConn
-	state              int64
+	state              common.ServiceState
 	wgServer           sync.WaitGroup
 	bulkInsert         int
 	bulkInsertDeadline time.Duration
@@ -224,7 +223,7 @@ func (s *FlowServer) storeFlows(flows []*flow.Flow) {
 
 // Start the flow server
 func (s *FlowServer) Start() {
-	atomic.StoreInt64(&s.state, common.RunningState)
+	s.state.Store(common.RunningState)
 	s.wgServer.Add(1)
 
 	s.conn.Serve(s.ch, s.quit, &s.wgServer)
@@ -257,7 +256,7 @@ func (s *FlowServer) Start() {
 
 // Stop the server
 func (s *FlowServer) Stop() {
-	if atomic.CompareAndSwapInt64(&s.state, common.RunningState, common.StoppingState) {
+	if s.state.CompareAndSwap(common.RunningState, common.StoppingState) {
 		s.quit <- struct{}{}
 		s.quit <- struct{}{}
 		s.wgServer.Wait()

flow/table.go

@@ -20,7 +20,6 @@ package flow
 import (
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/google/gopacket"
@@ -63,7 +62,7 @@ type Table struct {
 	flushDone         chan bool
 	query             chan *TableQuery
 	reply             chan []byte
-	state             int64
+	state             common.ServiceState
 	lockState         common.RWMutex
 	wg                sync.WaitGroup
 	quit              chan bool
@@ -348,7 +347,7 @@ func (ft *Table) Query(query *TableQuery) []byte {
 	ft.lockState.Lock()
 	defer ft.lockState.Unlock()
 
-	if atomic.LoadInt64(&ft.state) == common.RunningState {
+	if ft.state.Load() == common.RunningState {
 		ft.query <- query
 
 		timer := time.NewTicker(1 * time.Second)
@@ -359,7 +358,7 @@ func (ft *Table) Query(query *TableQuery) []byte {
 			case r := <-ft.reply:
 				return r
 			case <-timer.C:
-				if atomic.LoadInt64(&ft.state) != common.RunningState {
+				if ft.state.Load() != common.RunningState {
 					return nil
 				}
 			}
@@ -478,8 +477,8 @@ func (ft *Table) processEBPFFlow(ebpfFlow *EBPFFlow) {
 }
 
 // State returns the state of the flow table, stopped, running...
-func (ft *Table) State() int64 {
-	return atomic.LoadInt64(&ft.state)
+func (ft *Table) State() common.ServiceState {
+	return ft.state.Load()
 }
 
 // Run background jobs, like update/expire entries event
@@ -508,7 +507,7 @@ func (ft *Table) Run() {
 	ft.query = make(chan *TableQuery, 100)
 	ft.reply = make(chan []byte, 100)
 
-	atomic.StoreInt64(&ft.state, common.RunningState)
+	ft.state.Store(common.RunningState)
 	for {
 		select {
 		case <-ft.quit:
@@ -584,7 +583,7 @@ func (ft *Table) Stop() {
 	ft.lockState.Lock()
 	defer ft.lockState.Unlock()
 
-	if atomic.CompareAndSwapInt64(&ft.state, common.RunningState, common.StoppingState) {
+	if ft.state.CompareAndSwap(common.RunningState, common.StoppingState) {
 		ft.quit <- true
 		ft.wg.Wait()
 

probe/bundle.go

@@ -32,6 +32,16 @@ type Handler interface {
 	Stop()
 }
 
+// ServiceStatus describes the status returned by GetStatus
+type ServiceStatus struct {
+	Status common.ServiceState
+}
+
+// StatusReporter can be implemented by probes to report their status
+type StatusReporter interface {
+	GetStatus() interface{}
+}
+
 // Bundle describes a bundle of probes (topology of flow)
 type Bundle struct {
 	common.RWMutex

tests/scale_test.go

@@ -312,7 +312,7 @@ func TestScaleHA(t *testing.T) {
 	gh := gclient.NewGremlinQueryHelper(authOptions)
 
 	// expected 1 either Incomer or Outgoer
-	if err = common.Retry(func() error { return checkPeers(client, 1, common.RunningState) }, 5, time.Second); err != nil {
+	if err = common.Retry(func() error { return checkPeers(client, 1, websocket.ConnState(common.RunningState)) }, 5, time.Second); err != nil {
 		execCmds(t, tearDownCmds...)
 		t.Fatal(err)
 	}
@@ -405,7 +405,7 @@ func TestScaleHA(t *testing.T) {
 	}
 	execCmds(t, setupCmds...)
 
-	if err = checkPeers(client, 0, common.RunningState); err != nil {
+	if err = checkPeers(client, 0, websocket.ConnState(common.RunningState)); err != nil {
 		execCmds(t, tearDownCmds...)
 		t.Fatal(err)
 	}
@@ -453,7 +453,7 @@ func TestScaleHA(t *testing.T) {
 	execCmds(t, setupCmds...)
 
 	if err = common.Retry(func() error {
-		return checkPeers(client, 1, common.RunningState)
+		return checkPeers(client, 1, websocket.ConnState(common.RunningState))
 	}, 15, time.Second); err != nil {
 		execCmds(t, tearDownCmds...)
 		t.Fatal(err)