refactor DiameterAVP struct, improve its performance

MATRIXXSoftware · Nov 11, 2023 · 4413d3c · 4413d3c
1 parent 6e7870e
commit 4413d3c
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 82 deletions.
diff --git a/diameter.go b/diameter.go
@@ -23,18 +23,11 @@ type DiameterClient struct {
 }
 
 type DiameterMessage struct {
-	name string // test
-
+	name    string // not exactly useful
 	diamMsg *diam.Message
-	avps    []*DiameterAVP // may need to mutex lock this
 }
 
-type DiameterAVP struct {
-	code   uint32
-	flags  uint8
-	vendor uint32
-	data   datatype.Type
-}
+type DataType struct{}
 
 func (*Diameter) NewClient() (*DiameterClient, error) {
 
@@ -87,66 +80,6 @@ func handleCCA(hopIds map[uint32]chan *diam.Message) diam.HandlerFunc {
 	}
 }
 
-func (*Diameter) NewMessage(name string) *DiameterMessage {
-
-	diamMsg := diam.NewRequest(diam.CreditControl, 4, dict.Default)
-
-	return &DiameterMessage{
-		name:    name,
-		diamMsg: diamMsg,
-		avps:    []*DiameterAVP{},
-	}
-}
-
-func (m *DiameterMessage) AddAVP() *DiameterAVP {
-	// populate later
-	avp := DiameterAVP{
-		code:   0,
-		flags:  0,
-		vendor: 0,
-		data:   nil,
-	}
-	m.avps = append(m.avps, &avp)
-	return &avp
-}
-
-func (a *DiameterAVP) XCode(code uint32) *DiameterAVP {
-	a.code = code
-	return a
-}
-
-func (a *DiameterAVP) XMbit() *DiameterAVP {
-	a.flags = a.flags | avp.Mbit
-	return a
-}
-
-func (a *DiameterAVP) XPbit() *DiameterAVP {
-	a.flags = a.flags | avp.Pbit
-	return a
-}
-
-func (a *DiameterAVP) XVbit() *DiameterAVP {
-	a.flags = a.flags | avp.Vbit
-	return a
-}
-
-func (a *DiameterAVP) XVendor(vendor uint32) *DiameterAVP {
-	a.vendor = vendor
-	return a
-}
-
-func (a *DiameterAVP) XUTF8String(value string) *DiameterAVP {
-	a.data = datatype.UTF8String(value)
-	return a
-}
-
-func (a *DiameterAVP) XDiameterIdentity(value string) *DiameterAVP {
-	a.data = datatype.DiameterIdentity(value)
-	return a
-}
-
-// TODO add more data type
-
 func (c *DiameterClient) Connect(address string) error {
 	if c.conn != nil {
 		return nil
@@ -171,10 +104,6 @@ func (d *Diameter) Send(client *DiameterClient, msg *DiameterMessage) (uint32, e
 
 	req := msg.diamMsg
 
-	for _, avp := range msg.avps {
-		req.NewAVP(avp.code, avp.flags, avp.vendor, avp.data)
-	}
-
 	// Keep track of Hop-by-Hop ID
 	hopByHopID := req.Header.HopByHopID
 	client.hopIds[hopByHopID] = make(chan *diam.Message)
@@ -200,6 +129,97 @@ func (d *Diameter) Send(client *DiameterClient, msg *DiameterMessage) (uint32, e
 	return uint32(resultCode), nil
 }
 
+func (*Diameter) NewMessage(name string) *DiameterMessage {
+
+	diamMsg := diam.NewRequest(diam.CreditControl, 4, dict.Default)
+
+	return &DiameterMessage{
+		name:    name,
+		diamMsg: diamMsg,
+		//avps:    []*DiameterAVP{},
+	}
+}
+
+func (m *DiameterMessage) XAVP(code uint32, vendor uint32, flags uint8, data datatype.Type) {
+	m.diamMsg.NewAVP(code, flags, vendor, data)
+}
+
+func (*Diameter) XDataType() DataType {
+	return DataType{}
+}
+
+func (d *DataType) XAddress(value string) datatype.Type {
+	return datatype.Address(value)
+}
+
+func (d *DataType) XDiameterIdentity(value string) datatype.Type {
+	return datatype.DiameterIdentity(value)
+}
+
+func (d *DataType) XDiameterURI(value string) datatype.Type {
+	return datatype.DiameterURI(value)
+}
+
+func (d *DataType) XEnumerated(value int32) datatype.Type {
+	return datatype.Enumerated(value)
+}
+
+func (d *DataType) XFloat32(value float32) datatype.Type {
+	return datatype.Float32(value)
+}
+
+func (d *DataType) XFloat64(value float64) datatype.Type {
+	return datatype.Float64(value)
+}
+
+func (d *DataType) XGrouped(value string) datatype.Type {
+	return datatype.Grouped(value)
+}
+
+func (d *DataType) XIPFilterRule(value string) datatype.Type {
+	return datatype.IPFilterRule(value)
+}
+
+func (d *DataType) XIPv4(value string) datatype.Type {
+	return datatype.IPv4(value)
+}
+
+func (d *DataType) XIPv6(value string) datatype.Type {
+	return datatype.IPv6(value)
+}
+
+func (d *DataType) XInteger32(value int32) datatype.Type {
+	return datatype.Integer32(value)
+}
+
+func (d *DataType) XInteger64(value int64) datatype.Type {
+	return datatype.Integer64(value)
+}
+
+func (d *DataType) XOctetString(value string) datatype.Type {
+	return datatype.OctetString(value)
+}
+
+func (d *DataType) XQoSFilterRule(value string) datatype.Type {
+	return datatype.QoSFilterRule(value)
+}
+
+func (d *DataType) XTime(value time.Time) datatype.Type {
+	return datatype.Time(value)
+}
+
+func (d *DataType) XUTF8String(value string) datatype.Type {
+	return datatype.UTF8String(value)
+}
+
+func (d *DataType) XUnsigned32(value uint32) datatype.Type {
+	return datatype.Unsigned32(value)
+}
+
+func (d *DataType) XUnsigned64(value uint64) datatype.Type {
+	return datatype.Unsigned64(value)
+}
+
 func init() {
 	diameter := &Diameter{}
 	modules.Register("k6/x/diameter", diameter)

diff --git a/example.js b/example.js
@@ -2,25 +2,34 @@ import diameter from 'k6/x/diameter';
 import { check } from 'k6';
 
 export let options = {
-    iterations: 5,
-    vus: 2,
+    // iterations: 500000,
+    // vus: 32,
+
+    iterations: 2,
+    vus: 1,
 }
 
+let Vbit = 0x80;
+let Mbit = 0x40;
+let Pbit = 0x20;
 let client = diameter.newClient();
+let dataType = diameter.DataType();
 
 export default function () {
     client.connect("localhost:3868")
 
     let msg = diameter.newMessage("CCR");
-    msg.addAVP().Code(263).Mbit().UTF8String("Session-8888");      // Session ID
-    msg.addAVP().Code(264).DiameterIdentity("origin.host")         // Origin-Host
-    msg.addAVP().Code(296).DiameterIdentity("origin.realm")        // Origin-Realm
-    msg.addAVP().Code(283).DiameterIdentity("dest.host")           // Destination-Host
-    msg.addAVP().Code(293).DiameterIdentity("dest.realm")          // Destination-Realm
-    msg.addAVP().Code(1).Vendor(10415).UTF8String("ValueFooBar");  // User-Name
+    msg.AVP(264, 0,     0,    dataType.DiameterIdentity("origin.host"))     // Origin-Host
+    msg.AVP(296, 0,     0,    dataType.DiameterIdentity("origin.realm"))    // Origin-Realm
+    msg.AVP(283, 0,     0,    dataType.DiameterIdentity("dest.host"))       // Destination-Host
+    msg.AVP(293, 0,     0,    dataType.DiameterIdentity("dest.realm"))      // Destination-Realm
+
+    msg.AVP(263, 0,     Mbit, dataType.UTF8String("Session-8888"))          // Session ID
+    msg.AVP(416, 0,     Mbit, dataType.Enumerated(1))                       // CC-Request-Type
+    msg.AVP(415, 0,     Mbit, dataType.Unsigned32(1000))                    // CC-Request-Number
 
     const response = diameter.send(client, msg);
-    console.log("Result Code:", response);
+    //console.log("Result Code:", response);
 
     check(response, {'Result-Code == 2001': r => r == 2001,});
 }