diff --git a/cmd/server.go b/cmd/server.go index 999193ba0..54603806f 100644 --- a/cmd/server.go +++ b/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.") diff --git a/config.go b/config.go index 0711e8dba..3e58a2d13 100644 --- a/config.go +++ b/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 diff --git a/ctl/config.go b/ctl/config.go index 75cf6adf3..d084c1e2c 100644 --- a/ctl/config.go +++ b/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" diff --git a/executor.go b/executor.go index a16c3f139..39753787e 100644 --- a/executor.go +++ b/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{} diff --git a/executor_test.go b/executor_test.go index 1e0cd8343..a89911f38 100644 --- a/executor_test.go +++ b/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 diff --git a/handler.go b/handler.go index fb2a35a0f..3e837bef3 100644 --- a/handler.go +++ b/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. diff --git a/pilosa.go b/pilosa.go index dd7402677..b50ebbf78 100644 --- a/pilosa.go +++ b/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. diff --git a/pql/ast.go b/pql/ast.go index b83b35731..19b5381ed 100644 --- a/pql/ast.go +++ b/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)) diff --git a/server.go b/server.go index 4a5b4192c..c283b3d0e 100644 --- a/server.go +++ b/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 diff --git a/server/server.go b/server/server.go index 73522b0fa..b8475718a 100644 --- a/server/server.go +++ b/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 {