From e3d6f966572afb8aa7ed50b5ab0b59482ee428f0 Mon Sep 17 00:00:00 2001 From: Ben Johnson Date: Wed, 3 May 2017 08:53:20 -0600 Subject: [PATCH] Add max-writes-per-requests limit. A configurable limit has been added to restrict the number of mutating calls in a `pql.Query`. This is to prevents requests from timing out from large queries. The default is set to 5000 writes per request and is configurable through the configuration file and the command line flags. --- cmd/server.go | 1 + config.go | 10 +++++++++- ctl/config.go | 1 + executor.go | 8 ++++++++ executor_test.go | 11 +++++++++++ handler.go | 7 ++++++- pilosa.go | 1 + pql/ast.go | 12 ++++++++++++ server.go | 4 ++++ server/server.go | 3 +++ 10 files changed, 56 insertions(+), 2 deletions(-) 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 {