Merge pull request #1628 from orozery/objectstore_marshaller

add objectstore marshaller

Author: safchain

Diff for: contrib/objectstore/main.go

@@ -9,6 +9,7 @@ import (
 	"github.com/skydive-project/skydive/common"
 	"github.com/skydive-project/skydive/config"
 	"github.com/skydive-project/skydive/contrib/objectstore/subscriber"
+	"github.com/skydive-project/skydive/contrib/objectstore/subscriber/flowtransformer"
 	shttp "github.com/skydive-project/skydive/http"
 	"github.com/skydive-project/skydive/logging"
 	"github.com/skydive-project/skydive/websocket"
@@ -39,6 +40,13 @@ func main() {
 	subscriberUsername := cfg.GetString("subscriber_username")
 	subscriberPassword := cfg.GetString("subscriber_password")
 	maxSecondsPerStream := cfg.GetInt("max_seconds_per_stream")
+	flowTransformerName := cfg.GetString("flow_transformer")
+
+	flowTransformer, err := flowtransformer.New(flowTransformerName)
+	if err != nil {
+		logging.GetLogger().Errorf("Failed to initialize flow transformer: %s", err.Error())
+		os.Exit(1)
+	}
 
 	authOpts := &shttp.AuthenticationOpts{
 		Username: subscriberUsername,
@@ -58,7 +66,7 @@ func main() {
 	}
 	structClient := wsClient.UpgradeToStructSpeaker()
 
-	s := subscriber.New(endpoint, region, bucket, accessKey, secretKey, objectPrefix, maxSecondsPerStream)
+	s := subscriber.New(endpoint, region, bucket, accessKey, secretKey, objectPrefix, maxSecondsPerStream, flowTransformer)
 
 	// subscribe to the flow updates
 	structClient.AddStructMessageHandler(s, []string{"flow"})

Diff for: contrib/objectstore/skydive-objectstore.yml.default

@@ -8,3 +8,4 @@
 # subscriber_username:
 # subscriber_password:
 # max_seconds_per_stream: 86400
+# flow_transformer: custom1

Diff for: contrib/objectstore/subscriber/flowtransformer/custom1/custom1.go

@@ -0,0 +1,42 @@
+package custom1
+
+import (
+	"github.com/skydive-project/skydive/flow"
+)
+
+// Flow represents a transformed flow
+type Flow struct {
+	UUID             string
+	LayersPath       string
+	Network          *flow.FlowLayer
+	Transport        *flow.TransportLayer
+	LastUpdateMetric *flow.FlowMetric
+	Metric           *flow.FlowMetric
+	Start            int64
+	Last             int64
+	FinishType       flow.FlowFinishType
+}
+
+// FlowTransformer is a custom transformer for flows
+type FlowTransformer struct {
+}
+
+// Transform transforms a flow before being stored
+func (m *FlowTransformer) Transform(f *flow.Flow) interface{} {
+	return &Flow{
+		UUID:             f.UUID,
+		LayersPath:       f.LayersPath,
+		Network:          f.Network,
+		Transport:        f.Transport,
+		LastUpdateMetric: f.LastUpdateMetric,
+		Metric:           f.Metric,
+		Start:            f.Start,
+		Last:             f.Last,
+		FinishType:       f.FinishType,
+	}
+}
+
+// New returns a new Custom1Marshaller
+func New() *FlowTransformer {
+	return &FlowTransformer{}
+}

Diff for: contrib/objectstore/subscriber/flowtransformer/flowtransformer.go

@@ -0,0 +1,25 @@
+package flowtransformer
+
+import (
+	"fmt"
+	"github.com/skydive-project/skydive/contrib/objectstore/subscriber/flowtransformer/custom1"
+	"github.com/skydive-project/skydive/flow"
+)
+
+// FlowTransformer allows generic transformations of a flow
+type FlowTransformer interface {
+	// Transform transforms a flow before being stored
+	Transform(flow *flow.Flow) interface{}
+}
+
+// New creates a new flow transformer based on a name string
+func New(flowTransformerName string) (FlowTransformer, error) {
+	switch flowTransformerName {
+	case "custom1":
+		return custom1.New(), nil
+	case "":
+		return nil, nil
+	default:
+		return nil, fmt.Errorf("Marshaller '%s' is not supported", flowTransformerName)
+	}
+}

Diff for: contrib/objectstore/subscriber/objectstore.go

@@ -13,6 +13,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 
 	"github.com/skydive-project/skydive/contrib/objectstore/subscriber/client"
+	"github.com/skydive-project/skydive/contrib/objectstore/subscriber/flowtransformer"
 	"github.com/skydive-project/skydive/flow"
 	"github.com/skydive-project/skydive/logging"
 	ws "github.com/skydive-project/skydive/websocket"
@@ -33,6 +34,7 @@ type Subscriber struct {
 	currentStream     stream
 	maxStreamDuration time.Duration
 	objectStoreClient client.Client
+	flowTransformer   flowtransformer.FlowTransformer
 }
 
 // OnStructMessage is triggered when WS server sends us a message.
@@ -58,7 +60,21 @@ func (s *Subscriber) StoreFlows(flows []*flow.Flow) error {
 		return nil
 	}
 
-	flowsString, err := json.Marshal(flows)
+	var jsonMarshalInput interface{}
+	if s.flowTransformer == nil {
+		jsonMarshalInput = flows
+	} else {
+		transformedFlows := make([]interface{}, 0, len(flows))
+		for _, f := range flows {
+			transformedFlow := s.flowTransformer.Transform(f)
+			if transformedFlow != nil {
+				transformedFlows = append(transformedFlows, transformedFlow)
+			}
+		}
+		jsonMarshalInput = transformedFlows
+	}
+
+	flowsBytes, err := json.Marshal(jsonMarshalInput)
 	if err != nil {
 		logging.GetLogger().Error("Error encoding flows: ", err)
 		return err
@@ -83,7 +99,7 @@ func (s *Subscriber) StoreFlows(flows []*flow.Flow) error {
 	// gzip
 	var b bytes.Buffer
 	w := gzip.NewWriter(&b)
-	w.Write([]byte(flowsString))
+	w.Write(flowsBytes)
 	w.Close()
 
 	currentStream := s.currentStream
@@ -106,13 +122,14 @@ func (s *Subscriber) StoreFlows(flows []*flow.Flow) error {
 }
 
 // New returns a new flows subscriber writing to an object storage service
-func New(endpoint, region, bucket, accessKey, secretKey, objectPrefix string, maxSecondsPerStream int) *Subscriber {
+func New(endpoint, region, bucket, accessKey, secretKey, objectPrefix string, maxSecondsPerStream int, flowTransformer flowtransformer.FlowTransformer) *Subscriber {
 	objectStoreClient := client.New(endpoint, region, accessKey, secretKey)
 	s := &Subscriber{
 		bucket:            bucket,
 		objectPrefix:      objectPrefix,
 		maxStreamDuration: time.Second * time.Duration(maxSecondsPerStream),
 		objectStoreClient: objectStoreClient,
+		flowTransformer:   flowTransformer,
 	}
 	return s
 }