Move OnSend bucket increment to network/message/ package

api/send.go

@@ -50,7 +50,6 @@ func (c *Client) SendUnsafe(m message.Send, param params.Unsafe) ([]id.Round,
 func (c *Client) SendCMIX(msg format.Message, recipientID *id.ID,
 	param params.CMIX) (id.Round, ephemeral.Id, error) {
 	jww.INFO.Printf("SendCMIX(%s)", string(msg.GetContents()))
-	c.OnSend(1)
 	return c.network.SendCMIX(msg, recipientID, param)
 }
 
@@ -59,7 +58,6 @@ func (c *Client) SendCMIX(msg format.Message, recipientID *id.ID,
 // round ID of the round the payload was sent or an error if it fails.
 func (c *Client) SendManyCMIX(messages []message.TargetedCmixMessage,
 	params params.CMIX) (id.Round, []ephemeral.Id, error) {
-	c.OnSend(uint32(len(messages)))
 	return c.network.SendManyCMIX(messages, params)
 }
 
@@ -75,13 +73,3 @@ func (c *Client) NewCMIXMessage(contents []byte) (format.Message, error) {
 	msg.SetContents(contents)
 	return msg, nil
 }
-
-// OnSend performs a bucket addition on a call to Client.SendCMIX or
-// Client.SendManyCMIX, updating the bucket for the amount of messages sent.
-func (c *Client) OnSend(messages uint32) {
-	rateLimitingParam := c.storage.GetBucketParams().Get()
-	c.storage.GetBucket().AddWithExternalParams(messages,
-		rateLimitingParam.Capacity, rateLimitingParam.LeakedTokens,
-		rateLimitingParam.LeakDuration)
-
-}

network/message/sendCmix.go

@@ -187,6 +187,7 @@ func sendCmixHelper(sender *gateway.Sender, msg format.Message,
 				elapsed, numRoundTries)
 			jww.INFO.Print(m)
 			events.Report(1, "MessageSend", "Metric", m)
+			onSend(1, session)
 			return id.Round(bestRound.ID), ephID, nil
 		} else {
 			jww.FATAL.Panicf("Gateway %s returned no error, but failed "+
@@ -198,3 +199,13 @@ func sendCmixHelper(sender *gateway.Sender, msg format.Message,
 	return 0, ephemeral.Id{}, errors.New("failed to send the message, " +
 		"unknown error")
 }
+
+// OnSend performs a bucket addition on a call to Manager.SendCMIX or
+// Manager.SendManyCMIX, updating the bucket for the amount of messages sent.
+func onSend(messages uint32, session *storage.Session) {
+	rateLimitingParam := session.GetBucketParams().Get()
+	session.GetBucket().AddWithExternalParams(messages,
+		rateLimitingParam.Capacity, rateLimitingParam.LeakedTokens,
+		rateLimitingParam.LeakDuration)
+
+}

network/message/sendManyCmix.go

@@ -203,6 +203,7 @@ func sendManyCmixHelper(sender *gateway.Sender,
 				"in round %d", ephemeralIDsString, recipientString, bestRound.ID)
 			jww.INFO.Print(m)
 			events.Report(1, "MessageSendMany", "Metric", m)
+			onSend(uint32(len(msgs)), session)
 			return id.Round(bestRound.ID), ephemeralIDs, nil
 		} else {
 			jww.FATAL.Panicf("Gateway %s returned no error, but failed to "+