Merge pull request #14 from luissimas/feature/custom-metrics

Feature/custom metrics

README.md

@@ -43,7 +43,7 @@ export default function () {
     ])))
 
     const cca = client.send(ccr)
-    console.log("CCA: ", cca)
+    console.log(`CCA: ${cca}`)
 
     const resultCode = cca.findAVP(code.ResultCode, 0)
     check(resultCode, {'Result-Code == 2001': r => r == 2001,})

diameter.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"strconv"
 	"time"
 
 	"github.com/fiorix/go-diameter/v4/diam"
@@ -13,6 +14,8 @@ import (
 	"github.com/fiorix/go-diameter/v4/diam/dict"
 	"github.com/fiorix/go-diameter/v4/diam/sm"
 	log "github.com/sirupsen/logrus"
+	"go.k6.io/k6/js/modules"
+	"go.k6.io/k6/metrics"
 )
 
 type DiameterClient struct {
@@ -21,6 +24,8 @@ type DiameterClient struct {
 	hopIds            map[uint32]chan *diam.Message
 	requestTimeout    time.Duration
 	transportProtocol string
+	metrics           DiameterMetrics
+	vu                modules.VU
 }
 
 type DiameterMessage struct {
@@ -37,8 +42,7 @@ type GroupedAVP struct {
 
 type Dict struct{}
 
-func (*Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) {
-
+func (d *Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) {
 	config, err := parseConfig(arg)
 	if err != nil {
 		return nil, err
@@ -120,6 +124,8 @@ func (*Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) {
 		hopIds:            hopIds,
 		requestTimeout:    config.RequestTimeout.Duration,
 		transportProtocol: *config.TransportProtocol,
+		metrics:           d.metrics,
+		vu:                d.vu,
 	}, nil
 }
 
@@ -142,17 +148,16 @@ func (c *DiameterClient) Connect(address string) error {
 
 	conn, err := c.client.DialNetwork(c.transportProtocol, address)
 	if err != nil {
-		log.Errorf("Error connecting to %s, %v\n", "localhost:3368", err)
+		log.Errorf("Error connecting to %s, %v\n", address, err)
 		return err
 	}
-	log.Infof("Connected to %s\n", "localhost:3868")
+	log.Infof("Connected to %s\n", address)
 
 	c.conn = conn
 	return nil
 }
 
 func (c *DiameterClient) Send(msg *DiameterMessage) (*DiameterMessage, error) {
-
 	if c.conn == nil {
 		return nil, errors.New("Not connected")
 	}
@@ -166,23 +171,34 @@ func (c *DiameterClient) Send(msg *DiameterMessage) (*DiameterMessage, error) {
 	// Timeout settings
 	timeout := time.After(c.requestTimeout)
 
+	// Register current time to calculate request duration
+	sentAt := time.Now()
+	tags := map[string]string{
+		"cmd_code": strconv.FormatUint(uint64(msg.diamMsg.Header.CommandCode), 10),
+	}
+
 	// Send Request
 	_, err := req.WriteTo(c.conn)
 	if err != nil {
+		c.reportMetric(c.metrics.FailedRequestCount, time.Now(), 1, tags)
 		return nil, err
 	}
 
 	// Wait for Response
-	var resp *diam.Message
 	select {
-	case resp = <-c.hopIds[hopByHopID]:
+	case resp := <-c.hopIds[hopByHopID]:
+		now := time.Now()
+		c.reportMetric(c.metrics.RequestDuration, now, metrics.D(now.Sub(sentAt)), tags)
+		c.reportMetric(c.metrics.RequestCount, now, 1, tags)
+		c.reportMetric(c.metrics.FailedRequestCount, now, 0, tags)
+
+		delete(c.hopIds, hopByHopID)
+
+		return &DiameterMessage{diamMsg: resp}, nil
 	case <-timeout:
+		c.reportMetric(c.metrics.FailedRequestCount, time.Now(), 1, tags)
 		return nil, errors.New("Response timeout")
 	}
-
-	delete(c.hopIds, hopByHopID)
-
-	return &DiameterMessage{diamMsg: resp}, nil
 }
 
 func (*Diameter) NewMessage(cmd uint32, appid uint32) *DiameterMessage {

example/example.js

@@ -46,7 +46,7 @@ export default function () {
     ])))
 
     const cca = client.send(ccr)
-    console.log("CCA: ", cca)
+    console.log(`CCA: ${cca}`)
 
     const resultCode = cca.findAVP(code.ResultCode, 0)
     check(resultCode, {'Result-Code == 2001': r => r == 2001,})

go.mod

@@ -9,7 +9,10 @@ require (
 	github.com/fiorix/go-diameter/v4 v4.0.5-0.20231116194707-845be291bb4b
 )
 
-require golang.org/x/crypto v0.14.0 // indirect
+require (
+	golang.org/x/crypto v0.14.0 // indirect
+	google.golang.org/genproto v0.0.0-20231127180814-3a041ad873d4 // indirect
+)
 
 require (
 	github.com/dlclark/regexp2 v1.9.0 // indirect

go.sum

@@ -166,6 +166,8 @@ golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8T
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/genproto v0.0.0-20200903010400-9bfcb5116336 h1:ZcAny/XH59BbzUOKydQpvIlklwibW3T9SvDE5cGhdzc=
 google.golang.org/genproto v0.0.0-20200903010400-9bfcb5116336/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
+google.golang.org/genproto v0.0.0-20231127180814-3a041ad873d4 h1:W12Pwm4urIbRdGhMEg2NM9O3TWKjNcxQhs46V0ypf/k=
+google.golang.org/genproto v0.0.0-20231127180814-3a041ad873d4/go.mod h1:5RBcpGRxr25RbDzY5w+dmaqpSEvl8Gwl1x2CICf60ic=
 google.golang.org/grpc v1.41.0 h1:f+PlOh7QV4iIJkPrx5NQ7qaNGFQ3OTse67yaDHfju4E=
 google.golang.org/grpc v1.41.0/go.mod h1:U3l9uK9J0sini8mHphKoXyaqDA/8VyGnDee1zzIUK6k=
 google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=

metrics.go

@@ -0,0 +1,41 @@
+package diameter
+
+import (
+	"time"
+
+	"go.k6.io/k6/js/modules"
+	"go.k6.io/k6/metrics"
+)
+
+type DiameterMetrics struct {
+	RequestDuration    *metrics.Metric
+	RequestCount       *metrics.Metric
+	FailedRequestCount *metrics.Metric
+}
+
+func registerMetrics(vu modules.VU) DiameterMetrics {
+	registry := vu.InitEnv().Registry
+	metrics := DiameterMetrics{
+		RequestDuration:    registry.MustNewMetric("diameter_req_duration", metrics.Trend, metrics.Time),
+		RequestCount:       registry.MustNewMetric("diameter_req_count", metrics.Counter, metrics.Default),
+		FailedRequestCount: registry.MustNewMetric("diameter_req_failed", metrics.Rate, metrics.Default),
+	}
+	return metrics
+}
+
+func (c *DiameterClient) reportMetric(metric *metrics.Metric, now time.Time, value float64, tags map[string]string) {
+	state := c.vu.State()
+	ctx := c.vu.Context()
+	if state == nil || ctx == nil {
+		return
+	}
+
+	metrics.PushIfNotDone(ctx, state.Samples, metrics.Sample{
+		Time: now,
+		TimeSeries: metrics.TimeSeries{
+			Metric: metric,
+			Tags:   metrics.NewRegistry().RootTagSet().WithTagsFromMap(tags),
+		},
+		Value: value,
+	})
+}

module.go

@@ -1,23 +1,34 @@
 package diameter
 
 import (
-	"github.com/dop251/goja"
 	"go.k6.io/k6/js/modules"
 )
 
 type (
 	Diameter struct {
 		vu      modules.VU
-		exports *goja.Object
+		metrics DiameterMetrics
 	}
 	RootModule struct{}
-	Module     struct {
-		*Diameter
-	}
 )
 
 func init() {
-	modules.Register("k6/x/diameter", &Diameter{})
+	modules.Register("k6/x/diameter", New())
 	modules.Register("k6/x/diameter/avp", &AVP{})
 	modules.Register("k6/x/diameter/dict", &Dict{})
 }
+
+func New() *RootModule {
+	return &RootModule{}
+}
+
+func (*RootModule) NewModuleInstance(vu modules.VU) modules.Instance {
+	return &Diameter{
+		vu:      vu,
+		metrics: registerMetrics(vu),
+	}
+}
+
+func (d *Diameter) Exports() modules.Exports {
+	return modules.Exports{Default: d}
+}