cli/mrt: Drop queue size to have reasonable memory usage; make it configurable from cli

2 files changed, 14 insertions, 9 deletions

diff --git a/gobgp/cmd/common.go b/gobgp/cmd/common.go
index a0904c49..65a19a87 100644
--- a/gobgp/cmd/common.go
+++ b/gobgp/cmd/common.go
@@ -109,15 +109,16 @@ var actionOpts struct {
 }
 
 var mrtOpts struct {
-	OutputDir  string
-	FileFormat string
-	Filename   string `long:"filename" description:"MRT file name"`
+	OutputDir   string
+	FileFormat  string
+	Filename    string `long:"filename" description:"MRT file name"`
 	RecordCount int    `long:"count" description:"Number of records to inject"`
 	RecordSkip  int    `long:"skip" description:"Number of records to skip before injecting"`
-	Best       bool   `long:"only-best" description:"only keep best path routes"`
-	SkipV4     bool   `long:"no-ipv4" description:"Skip importing IPv4 routes"`
-	SkipV6     bool   `long:"no-ipv4" description:"Skip importing IPv6 routes"`
-	NextHop    net.IP `long:"nexthop" description:"Rewrite nexthop"`
+	QueueSize   int    `long:"batch-size" description:"Maximum number of updates to keep queued"`
+	Best        bool   `long:"only-best" description:"only keep best path routes"`
+	SkipV4      bool   `long:"no-ipv4" description:"Skip importing IPv4 routes"`
+	SkipV6      bool   `long:"no-ipv4" description:"Skip importing IPv6 routes"`
+	NextHop     net.IP `long:"nexthop" description:"Rewrite nexthop"`
 }
 
 func formatTimedelta(d int64) string {
diff --git a/gobgp/cmd/mrt.go b/gobgp/cmd/mrt.go
index ccd4a7f4..1f8d6694 100644
--- a/gobgp/cmd/mrt.go
+++ b/gobgp/cmd/mrt.go
@@ -39,8 +39,11 @@ func injectMrt() error {
 	}
 
 	idx := 0
+	if mrtOpts.QueueSize < 1 {
+		return fmt.Errorf("Specified queue size is smaller than 1, refusing to run with unbounded memory usage")
+	}
 
-	ch := make(chan []*table.Path, 1<<20)
+	ch := make(chan []*table.Path, mrtOpts.QueueSize)
 	go func() {
 
 		var peers []*mrt.Peer
@@ -138,7 +141,7 @@ func injectMrt() error {
 				}
 
 				idx += 1
-				if idx == mrtOpts.RecordCount + mrtOpts.RecordSkip {
+				if idx == mrtOpts.RecordCount+mrtOpts.RecordSkip {
 					break
 				}
 			}
@@ -209,6 +212,7 @@ func NewMrtCmd() *cobra.Command {
 	mrtCmd.PersistentFlags().BoolVarP(&mrtOpts.Best, "only-best", "", false, "inject only best paths")
 	mrtCmd.PersistentFlags().BoolVarP(&mrtOpts.SkipV4, "no-ipv4", "", false, "Do not import IPv4 routes")
 	mrtCmd.PersistentFlags().BoolVarP(&mrtOpts.SkipV6, "no-ipv6", "", false, "Do not import IPv6 routes")
+	mrtCmd.PersistentFlags().IntVarP(&mrtOpts.QueueSize, "queue-size", "", 1<<10, "Maximum number of updates to keep queued")
 	mrtCmd.PersistentFlags().IPVarP(&mrtOpts.NextHop, "nexthop", "", nil, "Overwrite nexthop")
 	return mrtCmd
 }