OAS: rework data feeding + sampleIDs (#917)

* Call OAS feeder

* Don't call old OAS code

* Rework calls

* Work on it

* Put back rules

* Make it compile

* start thinking of test

* Compiles

* Save

* Fixes

* Save

* Fixing

* Trying to fake conn

* add timeout

* Test timeout

* Fix tests

* Only build OAS for HTTP entries

* Remove some dead code

* Adding SampleIDs

* Cosmetics

* lint

* Revert rename

* Sample ID for content

* Cleanuo

* Add more sample IDs

* Checking hypothesis

* Move assignment place a bit

* Cosmetics

* Update test.yml

Co-authored-by: undera <undera@undera-old-desktop.home>
Co-authored-by: Igor Gov <iggvrv@gmail.com>

.github/workflows/test.yml

@@ -18,6 +18,7 @@ jobs:
   run-unit-tests:
     name: Unit Tests
     runs-on: ubuntu-latest
+    timeout-minutes: 10
     steps:
       - name: Check out code into the Go module directory
         uses: actions/checkout@v2

agent/main.go

@@ -371,6 +371,6 @@ func handleIncomingMessageAsTapper(socketConnection *websocket.Conn) {
 
 func initializeDependencies() {
 	dependency.RegisterGenerator(dependency.ServiceMapGeneratorDependency, func() interface{} { return servicemap.GetDefaultServiceMapInstance() })
-	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} { return oas.GetDefaultOasGeneratorInstance() })
+	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} { return oas.GetDefaultOasGeneratorInstance(nil) })
 	dependency.RegisterGenerator(dependency.EntriesProvider, func() interface{} { return &entries.BasenineEntriesProvider{} })
 }

agent/pkg/api/main.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/up9inc/mizu/agent/pkg/models"
 	"os"
 	"path"
 	"sort"
@@ -19,8 +20,6 @@ import (
 
 	"github.com/up9inc/mizu/agent/pkg/servicemap"
 
-	"github.com/up9inc/mizu/agent/pkg/models"
-	"github.com/up9inc/mizu/agent/pkg/oas"
 	"github.com/up9inc/mizu/agent/pkg/resolver"
 	"github.com/up9inc/mizu/agent/pkg/utils"
 
@@ -140,20 +139,6 @@ func startReadingChannel(outputItems <-chan *tapApi.OutputChannelItem, extension
 				rules, _, _ := models.RunValidationRulesState(*harEntry, mizuEntry.Destination.Name)
 				mizuEntry.Rules = rules
 			}
-
-			entryWSource := oas.EntryWithSource{
-				Entry:       *harEntry,
-				Source:      mizuEntry.Source.Name,
-				Destination: mizuEntry.Destination.Name,
-				Id:          mizuEntry.Id,
-			}
-
-			if entryWSource.Destination == "" {
-				entryWSource.Destination = mizuEntry.Destination.IP + ":" + mizuEntry.Destination.Port
-			}
-
-			oasGenerator := dependency.GetInstance(dependency.OasGeneratorDependency).(oas.OasGeneratorSink)
-			oasGenerator.PushEntry(&entryWSource)
 		}
 
 		data, err := json.Marshal(mizuEntry)

agent/pkg/controllers/oas_controller_test.go

@@ -1,8 +1,12 @@
 package controllers
 
 import (
+	"bytes"
+	basenine "github.com/up9inc/basenine/client/go"
+	"net"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/up9inc/mizu/agent/pkg/dependency"
 	"github.com/up9inc/mizu/agent/pkg/oas"
@@ -11,39 +15,55 @@ import (
 )
 
 func TestGetOASServers(t *testing.T) {
-	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} { return oas.GetDefaultOasGeneratorInstance() })
+	recorder, c := getRecorderAndContext()
-
-	recorder := httptest.NewRecorder()
-	c, _ := gin.CreateTestContext(recorder)
-	oas.GetDefaultOasGeneratorInstance().Start()
-	oas.GetDefaultOasGeneratorInstance().GetServiceSpecs().Store("some", oas.NewGen("some"))
 
 	GetOASServers(c)
 	t.Logf("Written body: %s", recorder.Body.String())
 }
 
 func TestGetOASAllSpecs(t *testing.T) {
-	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} { return oas.GetDefaultOasGeneratorInstance() })
+	recorder, c := getRecorderAndContext()
-
-	recorder := httptest.NewRecorder()
-	c, _ := gin.CreateTestContext(recorder)
-	oas.GetDefaultOasGeneratorInstance().Start()
-	oas.GetDefaultOasGeneratorInstance().GetServiceSpecs().Store("some", oas.NewGen("some"))
 
 	GetOASAllSpecs(c)
 	t.Logf("Written body: %s", recorder.Body.String())
 }
 
 func TestGetOASSpec(t *testing.T) {
-	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} { return oas.GetDefaultOasGeneratorInstance() })
+	recorder, c := getRecorderAndContext()
-
-	recorder := httptest.NewRecorder()
-	c, _ := gin.CreateTestContext(recorder)
-	oas.GetDefaultOasGeneratorInstance().Start()
-	oas.GetDefaultOasGeneratorInstance().GetServiceSpecs().Store("some", oas.NewGen("some"))
 
 	c.Params = []gin.Param{{Key: "id", Value: "some"}}
 
 	GetOASSpec(c)
 	t.Logf("Written body: %s", recorder.Body.String())
 }
+
+type fakeConn struct {
+	sendBuffer    *bytes.Buffer
+	receiveBuffer *bytes.Buffer
+}
+
+func (f fakeConn) Read(p []byte) (int, error)       { return f.sendBuffer.Read(p) }
+func (f fakeConn) Write(p []byte) (int, error)      { return f.receiveBuffer.Write(p) }
+func (fakeConn) Close() error                       { return nil }
+func (fakeConn) LocalAddr() net.Addr                { return nil }
+func (fakeConn) RemoteAddr() net.Addr               { return nil }
+func (fakeConn) SetDeadline(t time.Time) error      { return nil }
+func (fakeConn) SetReadDeadline(t time.Time) error  { return nil }
+func (fakeConn) SetWriteDeadline(t time.Time) error { return nil }
+
+func getRecorderAndContext() (*httptest.ResponseRecorder, *gin.Context) {
+	dummyConn := new(basenine.Connection)
+	dummyConn.Conn = fakeConn{
+		sendBuffer:    bytes.NewBufferString("\n"),
+		receiveBuffer: bytes.NewBufferString("\n"),
+	}
+	dependency.RegisterGenerator(dependency.OasGeneratorDependency, func() interface{} {
+		return oas.GetDefaultOasGeneratorInstance(dummyConn)
+	})
+
+	recorder := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(recorder)
+	oas.GetDefaultOasGeneratorInstance(dummyConn).Start()
+	oas.GetDefaultOasGeneratorInstance(dummyConn).GetServiceSpecs().Store("some", oas.NewGen("some"))
+	return recorder, c
+}

agent/pkg/oas/feeder_test.go

@@ -67,23 +67,23 @@ func fileSize(fname string) int64 {
 	return fi.Size()
 }
 
-func feedEntries(fromFiles []string, isSync bool) (count int, err error) {
+func feedEntries(fromFiles []string, isSync bool, gen *defaultOasGenerator) (count uint, err error) {
 	badFiles := make([]string, 0)
-	cnt := 0
+	cnt := uint(0)
 	for _, file := range fromFiles {
 		logger.Log.Info("Processing file: " + file)
 		ext := strings.ToLower(filepath.Ext(file))
-		eCnt := 0
+		eCnt := uint(0)
 		switch ext {
 		case ".har":
-			eCnt, err = feedFromHAR(file, isSync)
+			eCnt, err = feedFromHAR(file, isSync, gen)
 			if err != nil {
 				logger.Log.Warning("Failed processing file: " + err.Error())
 				badFiles = append(badFiles, file)
 				continue
 			}
 		case ".ldjson":
-			eCnt, err = feedFromLDJSON(file, isSync)
+			eCnt, err = feedFromLDJSON(file, isSync, gen)
 			if err != nil {
 				logger.Log.Warning("Failed processing file: " + err.Error())
 				badFiles = append(badFiles, file)
@@ -102,7 +102,7 @@ func feedEntries(fromFiles []string, isSync bool) (count int, err error) {
 	return cnt, nil
 }
 
-func feedFromHAR(file string, isSync bool) (int, error) {
+func feedFromHAR(file string, isSync bool, gen *defaultOasGenerator) (uint, error) {
 	fd, err := os.Open(file)
 	if err != nil {
 		panic(err)
@@ -121,16 +121,16 @@ func feedFromHAR(file string, isSync bool) (int, error) {
 		return 0, err
 	}
 
-	cnt := 0
+	cnt := uint(0)
 	for _, entry := range harDoc.Log.Entries {
 		cnt += 1
-		feedEntry(&entry, "", isSync, file)
+		feedEntry(&entry, "", file, gen, cnt)
 	}
 
 	return cnt, nil
 }
 
-func feedEntry(entry *har.Entry, source string, isSync bool, file string) {
+func feedEntry(entry *har.Entry, source string, file string, gen *defaultOasGenerator, cnt uint) {
 	entry.Comment = file
 	if entry.Response.Status == 302 {
 		logger.Log.Debugf("Dropped traffic entry due to permanent redirect status: %s", entry.StartedDateTime)
@@ -145,15 +145,11 @@ func feedEntry(entry *har.Entry, source string, isSync bool, file string) {
 		logger.Log.Errorf("Failed to parse entry URL: %v, err: %v", entry.Request.URL, err)
 	}
 
-	ews := EntryWithSource{Entry: *entry, Source: source, Destination: u.Host, Id: uint(0)}
+	ews := EntryWithSource{Entry: *entry, Source: source, Destination: u.Host, Id: cnt}
-	if isSync {
+	gen.handleHARWithSource(&ews)
-		GetDefaultOasGeneratorInstance().entriesChan <- ews // blocking variant, right?
-	} else {
-		GetDefaultOasGeneratorInstance().PushEntry(&ews)
-	}
 }
 
-func feedFromLDJSON(file string, isSync bool) (int, error) {
+func feedFromLDJSON(file string, isSync bool, gen *defaultOasGenerator) (uint, error) {
 	fd, err := os.Open(file)
 	if err != nil {
 		panic(err)
@@ -165,7 +161,7 @@ func feedFromLDJSON(file string, isSync bool) (int, error) {
 
 	var meta map[string]interface{}
 	buf := strings.Builder{}
-	cnt := 0
+	cnt := uint(0)
 	source := ""
 	for {
 		substr, isPrefix, err := reader.ReadLine()
@@ -196,7 +192,7 @@ func feedFromLDJSON(file string, isSync bool) (int, error) {
 				logger.Log.Warningf("Failed decoding entry: %s", line)
 			} else {
 				cnt += 1
-				feedEntry(&entry, source, isSync, file)
+				feedEntry(&entry, source, file, gen, cnt)
 			}
 		}
 	}

agent/pkg/oas/oas_generator.go

@@ -3,10 +3,13 @@ package oas
 import (
 	"context"
 	"encoding/json"
+	basenine "github.com/up9inc/basenine/client/go"
+	"github.com/up9inc/mizu/agent/pkg/har"
+	"github.com/up9inc/mizu/shared"
+	"github.com/up9inc/mizu/tap/api"
 	"net/url"
 	"sync"
 
-	"github.com/up9inc/mizu/agent/pkg/har"
 	"github.com/up9inc/mizu/shared/logger"
 )
 
@@ -15,10 +18,6 @@ var (
 	instance *defaultOasGenerator
 )
 
-type OasGeneratorSink interface {
-	PushEntry(entryWithSource *EntryWithSource)
-}
-
 type OasGenerator interface {
 	Start()
 	Stop()
@@ -32,12 +31,20 @@ type defaultOasGenerator struct {
 	ctx          context.Context
 	cancel       context.CancelFunc
 	serviceSpecs *sync.Map
-	entriesChan  chan EntryWithSource
+	dbConn       *basenine.Connection
 }
 
-func GetDefaultOasGeneratorInstance() *defaultOasGenerator {
+func GetDefaultOasGeneratorInstance(conn *basenine.Connection) *defaultOasGenerator {
 	syncOnce.Do(func() {
-		instance = NewDefaultOasGenerator()
+		if conn == nil {
+			c, err := basenine.NewConnection(shared.BasenineHost, shared.BaseninePort)
+			if err != nil {
+				panic(err)
+			}
+			conn = c
+		}
+
+		instance = NewDefaultOasGenerator(conn)
 		logger.Log.Debug("OAS Generator Initialized")
 	})
 	return instance
@@ -50,7 +57,6 @@ func (g *defaultOasGenerator) Start() {
 	ctx, cancel := context.WithCancel(context.Background())
 	g.cancel = cancel
 	g.ctx = ctx
-	g.entriesChan = make(chan EntryWithSource, 100) // buffer up to 100 entries for OAS processing
 	g.serviceSpecs = &sync.Map{}
 	g.started = true
 	go g.runGenerator()
@@ -70,80 +76,117 @@ func (g *defaultOasGenerator) IsStarted() bool {
 }
 
 func (g *defaultOasGenerator) runGenerator() {
+	// Make []byte channels to recieve the data and the meta
+	dataChan := make(chan []byte)
+	metaChan := make(chan []byte)
+
+	g.dbConn.Query("", dataChan, metaChan)
+
 	for {
 		select {
 		case <-g.ctx.Done():
 			logger.Log.Infof("OAS Generator was canceled")
 			return
 
-		case entryWithSource, ok := <-g.entriesChan:
+		case metaBytes, ok := <-metaChan:
+			if !ok {
+				logger.Log.Infof("OAS Generator - meta channel closed")
+				break
+			}
+			logger.Log.Debugf("Meta: %s", metaBytes)
+
+		case dataBytes, ok := <-dataChan:
 			if !ok {
 				logger.Log.Infof("OAS Generator - entries channel closed")
 				break
 			}
-			entry := entryWithSource.Entry
+
-			u, err := url.Parse(entry.Request.URL)
+			logger.Log.Debugf("Data: %s", dataBytes)
+			e := new(api.Entry)
+			err := json.Unmarshal(dataBytes, e)
 			if err != nil {
-				logger.Log.Errorf("Failed to parse entry URL: %v, err: %v", entry.Request.URL, err)
+				continue
+			}
+			g.handleEntry(e)
+		}
+	}
+}
+
+func (g *defaultOasGenerator) handleEntry(mizuEntry *api.Entry) {
+	if mizuEntry.Protocol.Name == "http" {
+		entry, err := har.NewEntry(mizuEntry.Request, mizuEntry.Response, mizuEntry.StartTime, mizuEntry.ElapsedTime)
+		if err != nil {
+			logger.Log.Warningf("Failed to turn MizuEntry %d into HAR Entry: %s", mizuEntry.Id, err)
+			return
 		}
 
-			val, found := g.serviceSpecs.Load(entryWithSource.Destination)
+		dest := mizuEntry.Destination.Name
-			var gen *SpecGen
+		if dest == "" {
-			if !found {
+			dest = mizuEntry.Destination.IP + ":" + mizuEntry.Destination.Port
-				gen = NewGen(u.Scheme + "://" + entryWithSource.Destination)
+		}
-				g.serviceSpecs.Store(entryWithSource.Destination, gen)
+
+		entryWSource := &EntryWithSource{
+			Entry:       *entry,
+			Source:      mizuEntry.Source.Name,
+			Destination: dest,
+			Id:          mizuEntry.Id,
+		}
+
+		g.handleHARWithSource(entryWSource)
 	} else {
-				gen = val.(*SpecGen)
+		logger.Log.Debugf("OAS: Unsupported protocol in entry %d: %s", mizuEntry.Id, mizuEntry.Protocol.Name)
 	}
+}
 
-			opId, err := gen.feedEntry(entryWithSource)
+func (g *defaultOasGenerator) handleHARWithSource(entryWSource *EntryWithSource) {
+	entry := entryWSource.Entry
+	gen := g.getGen(entryWSource.Destination, entry.Request.URL)
+
+	opId, err := gen.feedEntry(entryWSource)
 	if err != nil {
 		txt, suberr := json.Marshal(entry)
 		if suberr == nil {
 			logger.Log.Debugf("Problematic entry: %s", txt)
 		}
 
-				logger.Log.Warningf("Failed processing entry: %s", err)
+		logger.Log.Warningf("Failed processing entry %d: %s", entryWSource.Id, err)
-				continue
+		return
 	}
 
-			logger.Log.Debugf("Handled entry %s as opId: %s", entry.Request.URL, opId) // TODO: set opId back to entry?
+	logger.Log.Debugf("Handled entry %d as opId: %s", entryWSource.Id, opId) // TODO: set opId back to entry?
+}
+
+func (g *defaultOasGenerator) getGen(dest string, urlStr string) *SpecGen {
+	u, err := url.Parse(urlStr)
+	if err != nil {
+		logger.Log.Errorf("Failed to parse entry URL: %v, err: %v", urlStr, err)
 	}
+
+	val, found := g.serviceSpecs.Load(dest)
+	var gen *SpecGen
+	if !found {
+		gen = NewGen(u.Scheme + "://" + dest)
+		g.serviceSpecs.Store(dest, gen)
+	} else {
+		gen = val.(*SpecGen)
 	}
+	return gen
 }
 
 func (g *defaultOasGenerator) Reset() {
 	g.serviceSpecs = &sync.Map{}
 }
 
-func (g *defaultOasGenerator) PushEntry(entryWithSource *EntryWithSource) {
-	if !g.started {
-		return
-	}
-	select {
-	case g.entriesChan <- *entryWithSource:
-	default:
-		logger.Log.Warningf("OAS Generator - entry wasn't sent to channel because the channel has no buffer or there is no receiver")
-	}
-}
-
 func (g *defaultOasGenerator) GetServiceSpecs() *sync.Map {
 	return g.serviceSpecs
 }
 
-func NewDefaultOasGenerator() *defaultOasGenerator {
+func NewDefaultOasGenerator(c *basenine.Connection) *defaultOasGenerator {
 	return &defaultOasGenerator{
 		started:      false,
 		ctx:          nil,
 		cancel:       nil,
 		serviceSpecs: nil,
-		entriesChan:  nil,
+		dbConn:       c,
 	}
 }
-
-type EntryWithSource struct {
-	Source      string
-	Destination string
-	Entry       har.Entry
-	Id          uint
-}

agent/pkg/oas/oas_generator_test.go

@@ -0,0 +1,36 @@
+package oas
+
+import (
+	"encoding/json"
+	"github.com/up9inc/mizu/agent/pkg/har"
+	"sync"
+	"testing"
+)
+
+func TestOASGen(t *testing.T) {
+	gen := new(defaultOasGenerator)
+	gen.serviceSpecs = &sync.Map{}
+
+	e := new(har.Entry)
+	err := json.Unmarshal([]byte(`{"startedDateTime": "20000101","request": {"url": "https://host/path", "method": "GET"}, "response": {"status": 200}}`), e)
+	if err != nil {
+		panic(err)
+	}
+
+	ews := &EntryWithSource{
+		Destination: "some",
+		Entry:       *e,
+	}
+	gen.handleHARWithSource(ews)
+	g, ok := gen.serviceSpecs.Load("some")
+	if !ok {
+		panic("Failed")
+	}
+	sg := g.(*SpecGen)
+	spec, err := sg.GetSpec()
+	if err != nil {
+		panic(err)
+	}
+	specText, _ := json.Marshal(spec)
+	t.Log(string(specText))
+}

agent/pkg/oas/specgen.go

@@ -28,6 +28,13 @@ const CountersTotal = "x-counters-total"
 const CountersPerSource = "x-counters-per-source"
 const SampleId = "x-sample-entry"
 
+type EntryWithSource struct {
+	Source      string
+	Destination string
+	Entry       har.Entry
+	Id          uint
+}
+
 type reqResp struct { // hello, generics in Go
 	Req  *har.Request
 	Resp *har.Response
@@ -60,7 +67,7 @@ func (g *SpecGen) StartFromSpec(oas *openapi.OpenAPI) {
 	g.tree = new(Node)
 	for pathStr, pathObj := range oas.Paths.Items {
 		pathSplit := strings.Split(string(pathStr), "/")
-		g.tree.getOrSet(pathSplit, pathObj)
+		g.tree.getOrSet(pathSplit, pathObj, 0)
 
 		// clean "last entry timestamp" markers from the past
 		for _, pathAndOp := range g.tree.listOps() {
@@ -69,11 +76,11 @@ func (g *SpecGen) StartFromSpec(oas *openapi.OpenAPI) {
 	}
 }
 
-func (g *SpecGen) feedEntry(entryWithSource EntryWithSource) (string, error) {
+func (g *SpecGen) feedEntry(entryWithSource *EntryWithSource) (string, error) {
 	g.lock.Lock()
 	defer g.lock.Unlock()
 
-	opId, err := g.handlePathObj(&entryWithSource)
+	opId, err := g.handlePathObj(entryWithSource)
 	if err != nil {
 		return "", err
 	}
@@ -219,7 +226,7 @@ func (g *SpecGen) handlePathObj(entryWithSource *EntryWithSource) (string, error
 	} else {
 		split = strings.Split(urlParsed.Path, "/")
 	}
-	node := g.tree.getOrSet(split, new(openapi.PathObj))
+	node := g.tree.getOrSet(split, new(openapi.PathObj), entryWithSource.Id)
 	opObj, err := handleOpObj(entryWithSource, node.pathObj)
 
 	if opObj != nil {
@@ -242,12 +249,12 @@ func handleOpObj(entryWithSource *EntryWithSource, pathObj *openapi.PathObj) (*o
 		return nil, nil
 	}
 
-	err = handleRequest(&entry.Request, opObj, isSuccess)
+	err = handleRequest(&entry.Request, opObj, isSuccess, entryWithSource.Id)
 	if err != nil {
 		return nil, err
 	}
 
-	err = handleResponse(&entry.Response, opObj, isSuccess)
+	err = handleResponse(&entry.Response, opObj, isSuccess, entryWithSource.Id)
 	if err != nil {
 		return nil, err
 	}
@@ -257,6 +264,8 @@ func handleOpObj(entryWithSource *EntryWithSource, pathObj *openapi.PathObj) (*o
 		return nil, err
 	}
 
+	setSampleID(&opObj.Extensions, entryWithSource.Id)
+
 	return opObj, nil
 }
 
@@ -329,15 +338,10 @@ func handleCounters(opObj *openapi.Operation, success bool, entryWithSource *Ent
 		return err
 	}
 
-	err = opObj.Extensions.SetExtension(SampleId, entryWithSource.Id)
-	if err != nil {
-		return err
-	}
-
 	return nil
 }
 
-func handleRequest(req *har.Request, opObj *openapi.Operation, isSuccess bool) error {
+func handleRequest(req *har.Request, opObj *openapi.Operation, isSuccess bool, sampleId uint) error {
 	// TODO: we don't handle the situation when header/qstr param can be defined on pathObj level. Also the path param defined on opObj
 	urlParsed, err := url.Parse(req.URL)
 	if err != nil {
@@ -361,7 +365,7 @@ func handleRequest(req *har.Request, opObj *openapi.Operation, isSuccess bool) e
 		IsIgnored:      func(name string) bool { return false },
 		GeneralizeName: func(name string) string { return name },
 	}
-	handleNameVals(qstrGW, &opObj.Parameters, false)
+	handleNameVals(qstrGW, &opObj.Parameters, false, sampleId)
 
 	hdrGW := nvParams{
 		In:             openapi.InHeader,
@@ -369,7 +373,7 @@ func handleRequest(req *har.Request, opObj *openapi.Operation, isSuccess bool) e
 		IsIgnored:      isHeaderIgnored,
 		GeneralizeName: strings.ToLower,
 	}
-	handleNameVals(hdrGW, &opObj.Parameters, true)
+	handleNameVals(hdrGW, &opObj.Parameters, true, sampleId)
 
 	if isSuccess {
 		reqBody, err := getRequestBody(req, opObj)
@@ -378,12 +382,14 @@ func handleRequest(req *har.Request, opObj *openapi.Operation, isSuccess bool) e
 		}
 
 		if reqBody != nil {
+			setSampleID(&reqBody.Extensions, sampleId)
+
 			if req.PostData.Text == "" {
 				reqBody.Required = false
 			} else {
 
 				reqCtype, _ := getReqCtype(req)
-				reqMedia, err := fillContent(reqResp{Req: req}, reqBody.Content, reqCtype)
+				reqMedia, err := fillContent(reqResp{Req: req}, reqBody.Content, reqCtype, sampleId)
 				if err != nil {
 					return err
 				}
@@ -395,18 +401,20 @@ func handleRequest(req *har.Request, opObj *openapi.Operation, isSuccess bool) e
 	return nil
 }
 
-func handleResponse(resp *har.Response, opObj *openapi.Operation, isSuccess bool) error {
+func handleResponse(resp *har.Response, opObj *openapi.Operation, isSuccess bool, sampleId uint) error {
 	// TODO: we don't support "default" response
 	respObj, err := getResponseObj(resp, opObj, isSuccess)
 	if err != nil {
 		return err
 	}
 
-	handleRespHeaders(resp.Headers, respObj)
+	setSampleID(&respObj.Extensions, sampleId)
+
+	handleRespHeaders(resp.Headers, respObj, sampleId)
 
 	respCtype := getRespCtype(resp)
 	respContent := respObj.Content
-	respMedia, err := fillContent(reqResp{Resp: resp}, respContent, respCtype)
+	respMedia, err := fillContent(reqResp{Resp: resp}, respContent, respCtype, sampleId)
 	if err != nil {
 		return err
 	}
@@ -414,7 +422,7 @@ func handleResponse(resp *har.Response, opObj *openapi.Operation, isSuccess bool
 	return nil
 }
 
-func handleRespHeaders(reqHeaders []har.Header, respObj *openapi.ResponseObj) {
+func handleRespHeaders(reqHeaders []har.Header, respObj *openapi.ResponseObj, sampleId uint) {
 	visited := map[string]*openapi.HeaderObj{}
 	for _, pair := range reqHeaders {
 		if isHeaderIgnored(pair.Name) {
@@ -436,6 +444,8 @@ func handleRespHeaders(reqHeaders []har.Header, respObj *openapi.ResponseObj) {
 			logger.Log.Warningf("Failed to add example to a parameter: %s", err)
 		}
 		visited[nameGeneral] = param
+
+		setSampleID(&param.Extensions, sampleId)
 	}
 
 	// maintain "required" flag
@@ -456,13 +466,15 @@ func handleRespHeaders(reqHeaders []har.Header, respObj *openapi.ResponseObj) {
 	}
 }
 
-func fillContent(reqResp reqResp, respContent openapi.Content, ctype string) (*openapi.MediaType, error) {
+func fillContent(reqResp reqResp, respContent openapi.Content, ctype string, sampleId uint) (*openapi.MediaType, error) {
 	content, found := respContent[ctype]
 	if !found {
 		respContent[ctype] = &openapi.MediaType{}
 		content = respContent[ctype]
 	}
 
+	setSampleID(&content.Extensions, sampleId)
+
 	var text string
 	var isBinary bool
 	if reqResp.Req != nil {
@@ -474,10 +486,10 @@ func fillContent(reqResp reqResp, respContent openapi.Content, ctype string) (*o
 	if !isBinary && text != "" {
 		var exampleMsg []byte
 		// try treating it as json
-		any, isJSON := anyJSON(text)
+		anyVal, isJSON := anyJSON(text)
 		if isJSON {
 			// re-marshal with forced indent
-			if msg, err := json.MarshalIndent(any, "", "\t"); err != nil {
+			if msg, err := json.MarshalIndent(anyVal, "", "\t"); err != nil {
 				panic("Failed to re-marshal value, super-strange")
 			} else {
 				exampleMsg = msg

agent/pkg/oas/specgen_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"regexp"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -19,7 +20,7 @@ import (
 
 // if started via env, write file into subdir
 func outputSpec(label string, spec *openapi.OpenAPI, t *testing.T) string {
-	content, err := json.MarshalIndent(spec, "", "\t")
+	content, err := json.MarshalIndent(spec, "", "  ")
 	if err != nil {
 		panic(err)
 	}
@@ -48,14 +49,16 @@ func TestEntries(t *testing.T) {
 		t.Log(err)
 		t.FailNow()
 	}
-	GetDefaultOasGeneratorInstance().Start()
+
-	loadStartingOAS("test_artifacts/catalogue.json", "catalogue")
+	gen := NewDefaultOasGenerator(nil)
-	loadStartingOAS("test_artifacts/trcc.json", "trcc-api-service")
+	gen.serviceSpecs = new(sync.Map)
+	loadStartingOAS("test_artifacts/catalogue.json", "catalogue", gen.serviceSpecs)
+	loadStartingOAS("test_artifacts/trcc.json", "trcc-api-service", gen.serviceSpecs)
 
 	go func() {
 		for {
 			time.Sleep(1 * time.Second)
-			GetDefaultOasGeneratorInstance().GetServiceSpecs().Range(func(key, val interface{}) bool {
+			gen.serviceSpecs.Range(func(key, val interface{}) bool {
 				svc := key.(string)
 				t.Logf("Getting spec for %s", svc)
 				gen := val.(*SpecGen)
@@ -68,16 +71,14 @@ func TestEntries(t *testing.T) {
 		}
 	}()
 
-	cnt, err := feedEntries(files, true)
+	cnt, err := feedEntries(files, true, gen)
 	if err != nil {
 		t.Log(err)
 		t.Fail()
 	}
 
-	waitQueueProcessed()
-
 	svcs := strings.Builder{}
-	GetDefaultOasGeneratorInstance().GetServiceSpecs().Range(func(key, val interface{}) bool {
+	gen.serviceSpecs.Range(func(key, val interface{}) bool {
 		gen := val.(*SpecGen)
 		svc := key.(string)
 		svcs.WriteString(svc + ",")
@@ -99,7 +100,7 @@ func TestEntries(t *testing.T) {
 		return true
 	})
 
-	GetDefaultOasGeneratorInstance().GetServiceSpecs().Range(func(key, val interface{}) bool {
+	gen.serviceSpecs.Range(func(key, val interface{}) bool {
 		svc := key.(string)
 		gen := val.(*SpecGen)
 		spec, err := gen.GetSpec()
@@ -123,20 +124,18 @@ func TestEntries(t *testing.T) {
 }
 
 func TestFileSingle(t *testing.T) {
-	GetDefaultOasGeneratorInstance().Start()
+	gen := NewDefaultOasGenerator(nil)
-	GetDefaultOasGeneratorInstance().Reset()
+	gen.serviceSpecs = new(sync.Map)
 	// loadStartingOAS()
 	file := "test_artifacts/params.har"
 	files := []string{file}
-	cnt, err := feedEntries(files, true)
+	cnt, err := feedEntries(files, true, gen)
 	if err != nil {
 		logger.Log.Warning("Failed processing file: " + err.Error())
 		t.Fail()
 	}
 
-	waitQueueProcessed()
+	gen.serviceSpecs.Range(func(key, val interface{}) bool {
-
-	GetDefaultOasGeneratorInstance().GetServiceSpecs().Range(func(key, val interface{}) bool {
 		svc := key.(string)
 		gen := val.(*SpecGen)
 		spec, err := gen.GetSpec()
@@ -189,18 +188,7 @@ func TestFileSingle(t *testing.T) {
 	logger.Log.Infof("Processed entries: %d", cnt)
 }
 
-func waitQueueProcessed() {
+func loadStartingOAS(file string, label string, specs *sync.Map) {
-	for {
-		time.Sleep(100 * time.Millisecond)
-		queue := len(GetDefaultOasGeneratorInstance().entriesChan)
-		logger.Log.Infof("Queue: %d", queue)
-		if queue < 1 {
-			break
-		}
-	}
-}
-
-func loadStartingOAS(file string, label string) {
 	fd, err := os.Open(file)
 	if err != nil {
 		panic(err)
@@ -222,12 +210,14 @@ func loadStartingOAS(file string, label string) {
 	gen := NewGen(label)
 	gen.StartFromSpec(doc)
 
-	GetDefaultOasGeneratorInstance().GetServiceSpecs().Store(label, gen)
+	specs.Store(label, gen)
 }
 
 func TestEntriesNegative(t *testing.T) {
+	gen := NewDefaultOasGenerator(nil)
+	gen.serviceSpecs = new(sync.Map)
 	files := []string{"invalid"}
-	_, err := feedEntries(files, false)
+	_, err := feedEntries(files, false, gen)
 	if err == nil {
 		t.Logf("Should have failed")
 		t.Fail()
@@ -235,8 +225,10 @@ func TestEntriesNegative(t *testing.T) {
 }
 
 func TestEntriesPositive(t *testing.T) {
+	gen := NewDefaultOasGenerator(nil)
+	gen.serviceSpecs = new(sync.Map)
 	files := []string{"test_artifacts/params.har"}
-	_, err := feedEntries(files, false)
+	_, err := feedEntries(files, false, gen)
 	if err != nil {
 		t.Logf("Failed")
 		t.Fail()

agent/pkg/oas/test_artifacts/params.har.spec.json

@@ -21,9 +21,11 @@
             "description": "Successful call with status 200",
             "content": {
               "application/json": {
-                "example": null
+                "example": null,
-              }
+                "x-sample-entry": 4
               }
+            },
+            "x-sample-entry": 4
           }
         },
         "x-counters-per-source": {
@@ -45,7 +47,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750580.04,
-        "x-sample-entry": 0
+        "x-sample-entry": 4
       }
     },
     "/appears-twice": {
@@ -58,9 +60,11 @@
             "description": "Successful call with status 200",
             "content": {
               "application/json": {
-                "example": null
+                "example": null,
-              }
+                "x-sample-entry": 6
               }
+            },
+            "x-sample-entry": 6
           }
         },
         "x-counters-per-source": {
@@ -82,7 +86,7 @@
           "sumDuration": 1
         },
         "x-last-seen-ts": 1567750581.74,
-        "x-sample-entry": 0
+        "x-sample-entry": 6
       }
     },
     "/body-optional": {
@@ -94,8 +98,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 12
               }
+            },
+            "x-sample-entry": 12
           }
         },
         "x-counters-per-source": {
@@ -117,14 +124,16 @@
           "sumDuration": 0.01
         },
         "x-last-seen-ts": 1567750581.75,
-        "x-sample-entry": 0,
+        "x-sample-entry": 12,
         "requestBody": {
           "description": "Generic request body",
           "content": {
             "application/json": {
-              "example": "{\"key\", \"val\"}"
+              "example": "{\"key\", \"val\"}",
-            }
+              "x-sample-entry": 11
             }
+          },
+          "x-sample-entry": 12
         }
       }
     },
@@ -137,8 +146,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 13
               }
+            },
+            "x-sample-entry": 13
           }
         },
         "x-counters-per-source": {
@@ -160,15 +172,17 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750581.75,
-        "x-sample-entry": 0,
+        "x-sample-entry": 13,
         "requestBody": {
           "description": "Generic request body",
           "content": {
             "": {
-              "example": "body exists"
+              "example": "body exists",
+              "x-sample-entry": 13
             }
           },
-          "required": true
+          "required": true,
+          "x-sample-entry": 13
         }
       }
     },
@@ -182,9 +196,11 @@
             "description": "Successful call with status 200",
             "content": {
               "": {
-                "example": {}
+                "example": {},
-              }
+                "x-sample-entry": 9
               }
+            },
+            "x-sample-entry": 9
           }
         },
         "x-counters-per-source": {
@@ -206,7 +222,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750582.74,
-        "x-sample-entry": 0,
+        "x-sample-entry": 9,
         "requestBody": {
           "description": "Generic request body",
           "content": {
@@ -233,10 +249,12 @@
                   }
                 }
               },
-              "example": "--BOUNDARY\r\nContent-Disposition: form-data; name=\"file\"; filename=\"metadata.json\"\r\nContent-Type: application/json\r\n\r\n{\"functions\": 123}\r\n--BOUNDARY\r\nContent-Disposition: form-data; name=\"path\"\r\n\r\n/content/components\r\n--BOUNDARY--\r\n"
+              "example": "--BOUNDARY\r\nContent-Disposition: form-data; name=\"file\"; filename=\"metadata.json\"\r\nContent-Type: application/json\r\n\r\n{\"functions\": 123}\r\n--BOUNDARY\r\nContent-Disposition: form-data; name=\"path\"\r\n\r\n/content/components\r\n--BOUNDARY--\r\n",
+              "x-sample-entry": 9
             }
           },
-          "required": true
+          "required": true,
+          "x-sample-entry": 9
         }
       }
     },
@@ -249,8 +267,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 8
               }
+            },
+            "x-sample-entry": 8
           }
         },
         "x-counters-per-source": {
@@ -272,7 +293,7 @@
           "sumDuration": 1
         },
         "x-last-seen-ts": 1567750581.74,
-        "x-sample-entry": 0,
+        "x-sample-entry": 8,
         "requestBody": {
           "description": "Generic request body",
           "content": {
@@ -312,10 +333,12 @@
                   }
                 }
               },
-              "example": "agent-id=ade\u0026callback-url=\u0026token=sometoken"
+              "example": "agent-id=ade\u0026callback-url=\u0026token=sometoken",
+              "x-sample-entry": 8
             }
           },
-          "required": true
+          "required": true,
+          "x-sample-entry": 8
         }
       }
     },
@@ -331,8 +354,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 14
               }
+            },
+            "x-sample-entry": 14
           }
         },
         "x-counters-per-source": {
@@ -354,7 +380,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750582,
-        "x-sample-entry": 0
+        "x-sample-entry": 14
       },
       "parameters": [
         {
@@ -369,7 +395,8 @@
             "example #0": {
               "value": "234324"
             }
-          }
+          },
+          "x-sample-entry": 14
         }
       ]
     },
@@ -385,8 +412,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 18
               }
+            },
+            "x-sample-entry": 18
           }
         },
         "x-counters-per-source": {
@@ -408,7 +438,7 @@
           "sumDuration": 9.53e-7
         },
         "x-last-seen-ts": 1567750582.00,
-        "x-sample-entry": 0
+        "x-sample-entry": 18
       },
       "parameters": [
         {
@@ -436,7 +466,8 @@
             "example #4": {
               "value": "prefix-gibberish-afterwards"
             }
-          }
+          },
+          "x-sample-entry": 19
         }
       ]
     },
@@ -452,8 +483,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 15
               }
+            },
+            "x-sample-entry": 15
           }
         },
         "x-counters-per-source": {
@@ -475,7 +509,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750582.00,
-        "x-sample-entry": 0
+        "x-sample-entry": 15
       },
       "parameters": [
         {
@@ -503,7 +537,8 @@
             "example #4": {
               "value": "prefix-gibberish-afterwards"
             }
-          }
+          },
+          "x-sample-entry": 19
         }
       ]
     },
@@ -519,8 +554,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 16
               }
+            },
+            "x-sample-entry": 16
           }
         },
         "x-counters-per-source": {
@@ -542,7 +580,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750582.00,
-        "x-sample-entry": 0
+        "x-sample-entry": 16
       },
       "parameters": [
         {
@@ -570,7 +608,8 @@
             "example #4": {
               "value": "prefix-gibberish-afterwards"
             }
-          }
+          },
+          "x-sample-entry": 19
         }
       ]
     },
@@ -586,8 +625,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "": {}
+              "": {
+                "x-sample-entry": 19
               }
+            },
+            "x-sample-entry": 19
           }
         },
         "x-counters-per-source": {
@@ -609,7 +651,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750582.00,
-        "x-sample-entry": 0
+        "x-sample-entry": 19
       },
       "parameters": [
         {
@@ -624,7 +666,8 @@
             "example #0": {
               "value": "23421"
             }
-          }
+          },
+          "x-sample-entry": 19
         },
         {
           "name": "parampatternId",
@@ -651,7 +694,8 @@
             "example #4": {
               "value": "prefix-gibberish-afterwards"
             }
-          }
+          },
+          "x-sample-entry": 19
         }
       ]
     },
@@ -665,9 +709,11 @@
             "description": "Successful call with status 200",
             "content": {
               "application/json": {
-                "example": null
+                "example": null,
-              }
+                "x-sample-entry": 3
               }
+            },
+            "x-sample-entry": 3
           }
         },
         "x-counters-per-source": {
@@ -689,7 +735,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750579.74,
-        "x-sample-entry": 0
+        "x-sample-entry": 3
       },
       "parameters": [
         {
@@ -707,7 +753,8 @@
             "example #1": {
               "value": "<UUID4>"
             }
-          }
+          },
+          "x-sample-entry": 3
         }
       ]
     },
@@ -720,8 +767,11 @@
           "200": {
             "description": "Successful call with status 200",
             "content": {
-              "text/html": {}
+              "text/html": {
+                "x-sample-entry": 1
               }
+            },
+            "x-sample-entry": 1
           }
         },
         "x-counters-per-source": {
@@ -743,7 +793,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750483.86,
-        "x-sample-entry": 0
+        "x-sample-entry": 1
       },
       "parameters": [
         {
@@ -761,7 +811,8 @@
             "example #1": {
               "value": "<UUID4>"
             }
-          }
+          },
+          "x-sample-entry": 3
         }
       ]
     },
@@ -775,9 +826,11 @@
             "description": "Successful call with status 200",
             "content": {
               "application/json": {
-                "example": null
+                "example": null,
-              }
+                "x-sample-entry": 2
               }
+            },
+            "x-sample-entry": 2
           }
         },
         "x-counters-per-source": {
@@ -799,7 +852,7 @@
           "sumDuration": 0
         },
         "x-last-seen-ts": 1567750578.74,
-        "x-sample-entry": 0
+        "x-sample-entry": 2
       },
       "parameters": [
         {
@@ -817,7 +870,8 @@
             "example #1": {
               "value": "<UUID4>"
             }
-          }
+          },
+          "x-sample-entry": 3
         }
       ]
     }

agent/pkg/oas/tree.go

@@ -20,7 +20,7 @@ type Node struct {
 	children  []*Node
 }
 
-func (n *Node) getOrSet(path NodePath, existingPathObj *openapi.PathObj) (node *Node) {
+func (n *Node) getOrSet(path NodePath, existingPathObj *openapi.PathObj, sampleId uint) (node *Node) {
 	if existingPathObj == nil {
 		panic("Invalid function call")
 	}
@@ -70,6 +70,10 @@ func (n *Node) getOrSet(path NodePath, existingPathObj *openapi.PathObj) (node *
 		}
 	}
 
+	if node.pathParam != nil {
+		setSampleID(&node.pathParam.Extensions, sampleId)
+	}
+
 	// add example if it's a gibberish chunk
 	if node.pathParam != nil && !chunkIsParam {
 		exmp := &node.pathParam.Examples
@@ -85,7 +89,7 @@ func (n *Node) getOrSet(path NodePath, existingPathObj *openapi.PathObj) (node *
 
 	// TODO: eat up trailing slash, in a smart way: node.pathObj!=nil && path[1]==""
 	if len(path) > 1 {
-		return node.getOrSet(path[1:], existingPathObj)
+		return node.getOrSet(path[1:], existingPathObj, sampleId)
 	} else if node.pathObj == nil {
 		node.pathObj = existingPathObj
 	}

agent/pkg/oas/tree_test.go

@@ -20,10 +20,10 @@ func TestTree(t *testing.T) {
 	}
 
 	tree := new(Node)
-	for _, tc := range testCases {
+	for i, tc := range testCases {
 		split := strings.Split(tc.inp, "/")
 		pathObj := new(openapi.PathObj)
-		node := tree.getOrSet(split, pathObj)
+		node := tree.getOrSet(split, pathObj, uint(i))
 
 		fillPathParams(node, pathObj)
 

agent/pkg/oas/utils.go

@@ -115,7 +115,7 @@ type nvParams struct {
 	GeneralizeName func(name string) string
 }
 
-func handleNameVals(gw nvParams, params **openapi.ParameterList, checkIgnore bool) {
+func handleNameVals(gw nvParams, params **openapi.ParameterList, checkIgnore bool, sampleId uint) {
 	visited := map[string]*openapi.ParameterObj{}
 	for _, pair := range gw.Pairs {
 		if (checkIgnore && gw.IsIgnored(pair.Name)) || pair.Name == "" {
@@ -137,6 +137,8 @@ func handleNameVals(gw nvParams, params **openapi.ParameterList, checkIgnore boo
 			logger.Log.Warningf("Failed to add example to a parameter: %s", err)
 		}
 		visited[nameGeneral] = param
+
+		setSampleID(&param.Extensions, sampleId)
 	}
 
 	// maintain "required" flag
@@ -474,3 +476,15 @@ func intersectSliceWithMap(required []string, names map[string]struct{}) []strin
 	}
 	return required
 }
+
+func setSampleID(extensions *openapi.Extensions, id uint) {
+	if id > 0 {
+		if *extensions == nil {
+			*extensions = openapi.Extensions{}
+		}
+		err := (extensions).SetExtension(SampleId, id)
+		if err != nil {
+			logger.Log.Warningf("Failed to set sample ID: %s", err)
+		}
+	}
+}