Merge pull request #525 from benbjohnson/write-limit

Add max-writes-per-requests limit.

cmd/server.go

@@ -89,6 +89,7 @@ on the configured port.`,
 
 	flags.StringVarP(&Server.Config.DataDir, "data-dir", "d", "~/.pilosa", "Directory to store pilosa data files.")
 	flags.StringVarP(&Server.Config.Host, "bind", "b", ":10101", "Default URI on which pilosa should listen.")
+	flags.IntVarP(&Server.Config.MaxWritesPerRequest, "max-writes-per-request", "", Server.Config.MaxWritesPerRequest, "Number of write commands per request.")
 	flags.IntVarP(&Server.Config.Cluster.ReplicaN, "cluster.replicas", "", 1, "Number of hosts each piece of data should be stored on.")
 	flags.StringSliceVarP(&Server.Config.Cluster.Hosts, "cluster.hosts", "", []string{}, "Comma separated list of hosts in cluster.")
 	flags.StringSliceVarP(&Server.Config.Cluster.InternalHosts, "cluster.internal-hosts", "", []string{}, "Comma separated list of hosts in cluster used for internal communication.")

config.go

@@ -28,6 +28,9 @@ const (
 
 	// DefaultInternalPort the port the nodes intercommunicate on.
 	DefaultInternalPort = "14000"
+
+	// DefaultMaxWritesPerRequest is the default number of writes per request.
+	DefaultMaxWritesPerRequest = 5000
 )
 
 // Config represents the configuration for the command.
@@ -53,13 +56,18 @@ type Config struct {
 		Interval Duration `toml:"interval"`
 	} `toml:"anti-entropy"`
 
+	// Limits the number of mutating commands that can be in a single request to
+	// the server. This includes SetBit, ClearBit, SetRowAttrs & SetColumnAttrs.
+	MaxWritesPerRequest int `toml:"max-writes-per-request"`
+
 	LogPath string `toml:"log-path"`
 }
 
 // NewConfig returns an instance of Config with default options.
 func NewConfig() *Config {
 	c := &Config{
-		Host: DefaultHost + ":" + DefaultPort,
+		Host:                DefaultHost + ":" + DefaultPort,
+		MaxWritesPerRequest: DefaultMaxWritesPerRequest,
 	}
 	c.Cluster.ReplicaN = DefaultReplicaN
 	c.Cluster.Type = DefaultClusterType

ctl/config.go

@@ -40,6 +40,7 @@ func (cmd *ConfigCommand) Run(ctx context.Context) error {
 	fmt.Fprintln(cmd.Stdout, strings.TrimSpace(`
 data-dir = "~/.pilosa"
 bind = "localhost:10101"
+max-writes-per-request = 5000
 
 [cluster]
   poll-interval = "2m0s"

executor.go

@@ -49,6 +49,9 @@ type Executor struct {
 
 	// Client used for remote HTTP requests.
 	HTTPClient *http.Client
+
+	// Maximum number of SetBit() or ClearBit() commands per request.
+	MaxWritesPerRequest int
 }
 
 // NewExecutor returns a new instance of Executor.
@@ -65,6 +68,11 @@ func (e *Executor) Execute(ctx context.Context, index string, q *pql.Query, slic
 		return nil, ErrIndexRequired
 	}
 
+	// Verify that the number of writes do not exceed the maximum.
+	if e.MaxWritesPerRequest > 0 && q.WriteCallN() > e.MaxWritesPerRequest {
+		return nil, ErrTooManyWrites
+	}
+
 	// Default options.
 	if opt == nil {
 		opt = &ExecOptions{}

executor_test.go

@@ -699,6 +699,17 @@ func TestExecutor_Execute_Remote_TopN(t *testing.T) {
 	}
 }
 
+// Ensure executor returns an error if too many writes are in a single request.
+func TestExecutor_Execute_ErrMaxWritesPerRequest(t *testing.T) {
+	hldr := MustOpenHolder()
+	defer hldr.Close()
+	e := NewExecutor(hldr.Holder, NewCluster(1))
+	e.MaxWritesPerRequest = 3
+	if _, err := e.Execute(context.Background(), "i", MustParse(`SetBit() ClearBit() SetBit() SetBit()`), nil, nil); err != pilosa.ErrTooManyWrites {
+		t.Fatalf("unexpected error: %s", err)
+	}
+}
+
 // Executor represents a test wrapper for pilosa.Executor.
 type Executor struct {
 	*pilosa.Executor

handler.go

@@ -228,7 +228,12 @@ func (h *Handler) handlePostQuery(w http.ResponseWriter, r *http.Request) {
 
 	// Set appropriate status code, if there is an error.
 	if resp.Err != nil {
-		w.WriteHeader(http.StatusInternalServerError)
+		switch resp.Err {
+		case ErrTooManyWrites:
+			w.WriteHeader(http.StatusRequestEntityTooLarge)
+		default:
+			w.WriteHeader(http.StatusInternalServerError)
+		}
 	}
 
 	// Write response back to client.

pilosa.go

@@ -44,6 +44,7 @@ var (
 	// ErrFragmentNotFound is returned when a fragment does not exist.
 	ErrFragmentNotFound = errors.New("fragment not found")
 	ErrQueryRequired    = errors.New("query required")
+	ErrTooManyWrites    = errors.New("too many write commands")
 )
 
 // Regular expression to validate index and frame names.

pql/ast.go

@@ -28,6 +28,18 @@ type Query struct {
 	Calls []*Call
 }
 
+// WriteCallN returns the number of mutating calls.
+func (q *Query) WriteCallN() int {
+	var n int
+	for _, call := range q.Calls {
+		switch call.Name {
+		case "SetBit", "ClearBit", "SetRowAttrs", "SetColumnAttrs":
+			n++
+		}
+	}
+	return n
+}
+
 // String returns a string representation of the query.
 func (q *Query) String() string {
 	a := make([]string, len(q.Calls))

server.go

@@ -61,6 +61,9 @@ type Server struct {
 	AntiEntropyInterval time.Duration
 	PollingInterval     time.Duration
 
+	// Misc options.
+	MaxWritesPerRequest int
+
 	LogOutput io.Writer
 }
 
@@ -129,6 +132,7 @@ func (s *Server) Open() error {
 	e.Holder = s.Holder
 	e.Host = s.Host
 	e.Cluster = s.Cluster
+	e.MaxWritesPerRequest = s.MaxWritesPerRequest
 
 	// Initialize HTTP handler.
 	s.Handler.Broadcaster = s.Broadcaster

server/server.go

@@ -135,6 +135,9 @@ func (m *Command) SetupServer() error {
 	m.Server.Holder.Path = m.Config.DataDir
 	m.Server.Holder.Stats = pilosa.NewExpvarStatsClient()
 
+	// Copy configuration flags.
+	m.Server.MaxWritesPerRequest = m.Config.MaxWritesPerRequest
+
 	var err error
 	m.Server.Host, err = normalizeHost(m.Config.Host)
 	if err != nil {