diff --git a/api.go b/api.go index 6a802526b..913868414 100644 --- a/api.go +++ b/api.go @@ -25,6 +25,7 @@ import ( "io/ioutil" "strconv" "strings" + "sync" "time" "github.com/pilosa/pilosa/pql" @@ -42,6 +43,10 @@ type API struct { cluster *cluster server *Server + importWorkersWG sync.WaitGroup + importWorkerPoolSize int + importWork chan importJob + Serializer Serializer } @@ -58,9 +63,18 @@ func OptAPIServer(s *Server) apiOption { } } +func OptAPIImportWorkerPoolSize(size int) apiOption { + return func(a *API) error { + a.importWorkerPoolSize = size + return nil + } +} + // NewAPI returns a new API instance. func NewAPI(opts ...apiOption) (*API, error) { - api := &API{} + api := &API{ + importWorkerPoolSize: 2, + } for _, opt := range opts { err := opt(api) @@ -68,6 +82,16 @@ func NewAPI(opts ...apiOption) (*API, error) { return nil, errors.Wrap(err, "applying option") } } + + api.importWork = make(chan importJob, api.importWorkerPoolSize) + for i := 0; i < api.importWorkerPoolSize; i++ { + api.importWorkersWG.Add(1) + go func() { + importWorker(api.importWork) + defer api.importWorkersWG.Done() + }() + } + return api, nil } @@ -99,6 +123,13 @@ func (api *API) validate(f apiMethod) error { return newAPIMethodNotAllowedError(errors.Errorf("api method %s not allowed in state %s", f, state)) } +// Close closes the api and waits for it to shutdown. +func (api *API) Close() error { + close(api.importWork) + api.importWorkersWG.Wait() + return nil +} + // Query parses a PQL query out of the request and executes it. func (api *API) Query(ctx context.Context, req *QueryRequest) (QueryResponse, error) { span, ctx := tracing.StartSpanFromContext(ctx, "API.Query") @@ -271,6 +302,51 @@ func setUpImportOptions(opts ...ImportOption) (*ImportOptions, error) { return options, nil } +type importJob struct { + ctx context.Context + req *ImportRoaringRequest + shard uint64 + field *Field + errChan chan error +} + +func importWorker(importWork chan importJob) { + for j := range importWork { + err := func() error { + for viewName, viewData := range j.req.Views { + if viewName == "" { + viewName = viewStandard + } else { + viewName = fmt.Sprintf("%s_%s", viewStandard, viewName) + } + if len(viewData) == 0 { + return fmt.Errorf("no data to import for view: %s", viewName) + } + fileMagic := uint32(binary.LittleEndian.Uint16(viewData[0:2])) + if fileMagic == roaring.MagicNumber { // if pilosa roaring format + if err := j.field.importRoaring(j.ctx, viewData, j.shard, viewName, j.req.Clear); err != nil { + return errors.Wrap(err, "importing pilosa roaring") + } + } else { + // must make a copy of data to operate on locally on standard roaring format. + // field.importRoaring changes the standard roaring run format to pilosa roaring + data := make([]byte, len(viewData)) + copy(data, viewData) + if err := j.field.importRoaring(j.ctx, data, j.shard, viewName, j.req.Clear); err != nil { + return errors.Wrap(err, "importing standard roaring") + } + } + } + return nil + }() + + select { + case <-j.ctx.Done(): + case j.errChan <- err: + } + } +} + // ImportRoaring is a low level interface for importing data to Pilosa when // extremely high throughput is desired. The data must be encoded in a // particular way which may be unintuitive (discussed below). The data is merged @@ -298,7 +374,6 @@ func (api *API) ImportRoaring(ctx context.Context, indexName, fieldName string, } nodes := api.cluster.shardNodes(indexName, shard) - var eg errgroup.Group field := api.holder.Field(indexName, fieldName) if field == nil { @@ -310,48 +385,45 @@ func (api *API) ImportRoaring(ctx context.Context, indexName, fieldName string, return NewBadRequestError(errors.New("roaring import is only supported for set and time fields")) } + errCh := make(chan error, len(nodes)) + for _, node := range nodes { node := node if node.ID == api.server.nodeID { - eg.Go(func() error { - var err error - for viewName, viewData := range req.Views { - if viewName == "" { - viewName = viewStandard - } else { - viewName = fmt.Sprintf("%s_%s", viewStandard, viewName) - } - if len(viewData) == 0 { - return fmt.Errorf("no data to import for view: %s", viewName) - } - fileMagic := uint32(binary.LittleEndian.Uint16(viewData[0:2])) - if fileMagic == roaring.MagicNumber { // if pilosa roaring format - err = field.importRoaring(ctx, viewData, shard, viewName, req.Clear) - if err != nil { - return errors.Wrap(err, "importing pilosa roaring") - } - - } else { - // must make a copy of data to operate on locally on standard roaring format. - // field.importRoaring changes the standard roaring run format to pilosa roaring - data := make([]byte, len(viewData)) - copy(data, viewData) - err = field.importRoaring(ctx, data, shard, viewName, req.Clear) - if err != nil { - return errors.Wrap(err, "importing standard roaring") - } - } - } - return err - }) + api.importWork <- importJob{ + ctx: ctx, + req: req, + shard: shard, + field: field, + errChan: errCh, + } } else if !remote { // if remote == true we don't forward to other nodes // forward it on - eg.Go(func() error { - return api.server.defaultClient.ImportRoaring(ctx, &node.URI, indexName, fieldName, shard, true, req) - }) + go func() { + errCh <- api.server.defaultClient.ImportRoaring(ctx, &node.URI, indexName, fieldName, shard, true, req) + }() + } else { + errCh <- nil + } + } + + var maxNode int + for { + select { + case <-ctx.Done(): + return ctx.Err() + case nodeErr := <-errCh: + if nodeErr != nil { + return nodeErr + } + maxNode++ + } + + // Exit once all nodes are processed. + if maxNode == len(nodes) { + return nil } } - return eg.Wait() } // DeleteField removes the named field from the named index. If the index is not diff --git a/executor.go b/executor.go index 2110a92c6..4d1d3e4c9 100644 --- a/executor.go +++ b/executor.go @@ -1512,14 +1512,21 @@ func (e *executor) executeRowShard(ctx context.Context, index string, c *pql.Cal } // Union bitmaps across all time-based views. - row := &Row{} - for _, view := range viewsByTimeRange(viewStandard, fromTime, toTime, q) { + views := viewsByTimeRange(viewStandard, fromTime, toTime, q) + rows := make([]*Row, 0, len(views)) + for _, view := range views { f := e.Holder.fragment(index, fieldName, view, shard) if f == nil { continue } - row = row.Union(f.row(rowID)) + rows = append(rows, f.row(rowID)) } + if len(rows) == 0 { + return &Row{}, nil + } else if len(rows) == 1 { + return rows[0], nil + } + row := rows[0].Union(rows[1:]...) f.Stats.Count("range", 1, 1.0) return row, nil diff --git a/executor_test.go b/executor_test.go index be10b80bb..76a8300bd 100644 --- a/executor_test.go +++ b/executor_test.go @@ -2920,6 +2920,7 @@ func TestExecutor_Execute_ClearRow(t *testing.T) { Set(2, f=10, 2001-01-01T00:00)` readQueries := []string{ `Row(f=1, from=1999-12-31T00:00, to=2003-01-01T03:00)`, + `Row(f=1, from=2002-01-01T00:00, to=2002-01-02T00:00)`, `ClearRow(f=1)`, `Row(f=1, from=1999-12-31T00:00, to=2003-01-01T03:00)`, `Row(f=10, from=1999-12-31T00:00, to=2003-01-01T03:00)`, @@ -2931,20 +2932,26 @@ func TestExecutor_Execute_ClearRow(t *testing.T) { t.Fatalf("unexpected columns: %+v", columns) } + // Single day query (regression test) + if columns := responses[1].Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{7}) { + t.Fatalf("unexpected columns: %+v", columns) + } + // Clear the row and ensure we get a `true` response. - if res := responses[1].Results[0].(bool); !res { + if res := responses[2].Results[0].(bool); !res { t.Fatalf("unexpected clear row result: %+v", res) } // Ensure the row is empty. - if columns := responses[2].Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{}) { + if columns := responses[3].Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{}) { t.Fatalf("unexpected columns: %+v", columns) } // Ensure other rows were not affected. - if columns := responses[3].Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{2}) { + if columns := responses[4].Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{2}) { t.Fatalf("unexpected columns: %+v", columns) } + }) t.Run("Int", func(t *testing.T) { @@ -3292,6 +3299,7 @@ func TestExecutor_Execute_RowsTime(t *testing.T) { `Rows(f)`, `Rows(f, from=2002-01-01T00:00)`, `Rows(f, to=2003-02-03T00:00)`, + `Rows(f, from=2002-01-01T00:00, to=2002-01-02T00:00)`, } expResults := [][]uint64{ {1}, @@ -3300,6 +3308,7 @@ func TestExecutor_Execute_RowsTime(t *testing.T) { {1, 2, 3, 4, 13}, {2, 3, 4, 13}, {1, 2, 3, 13}, + {2}, } responses := runCallTest(t, writeQuery, readQueries, diff --git a/field.go b/field.go index b34628398..39c6c5363 100644 --- a/field.go +++ b/field.go @@ -144,6 +144,7 @@ func OptFieldTypeInt(min, max int64) FieldOption { fo.Type = FieldTypeInt fo.Min = min fo.Max = max + fo.Base = bsiBase(min, max) return nil } } @@ -518,7 +519,7 @@ func (f *Field) loadMeta() error { // Initialize "base" to "min" when upgrading from v1 BSI format. if pb.BitDepth == 0 { - pb.Base = pb.Min + pb.Base = bsiBase(pb.Min, pb.Max) pb.BitDepth = uint64(bitDepthInt64(pb.Max - pb.Min)) if pb.BitDepth == 0 { pb.BitDepth = 1 @@ -1520,6 +1521,18 @@ func isValidBSIGroupType(v string) bool { } } +// bsiBase is a helper function used to determine the default value +// for base. Because base is not exposed as a field option argument, +// it defaults to min, max, or 0 depending on the min/max range. +func bsiBase(min, max int64) int64 { + if min > 0 { + return min + } else if max < 0 { + return max + } + return 0 +} + // bsiGroup represents a group of range-encoded rows on a field. type bsiGroup struct { Name string `json:"name,omitempty"` diff --git a/field_internal_test.go b/field_internal_test.go index 13851a493..1ce7c8237 100644 --- a/field_internal_test.go +++ b/field_internal_test.go @@ -412,3 +412,29 @@ func TestField_PersistAvailableShardsFootprint(t *testing.T) { } } + +// Ensure that FieldOptions.Base defaults to the correct value. +func TestBSIGroup_BaseDefaultValue(t *testing.T) { + for i, tt := range []struct { + min int64 + max int64 + expBase int64 + }{ + {100, 200, 100}, + {-100, 100, 0}, + {-200, -100, -100}, + } { + fn := OptFieldTypeInt(tt.min, tt.max) + + // Apply functional option. + fo := FieldOptions{} + err := fn(&fo) + if err != nil { + t.Fatalf("test %d, applying functional option: %s", i, err.Error()) + } + + if fo.Base != tt.expBase { + t.Fatalf("test %d, unexpected FieldOptions.Base value. expected: %d, but got: %d", i, tt.expBase, fo.Base) + } + } +} diff --git a/row.go b/row.go index 79c918311..09a3390dc 100644 --- a/row.go +++ b/row.go @@ -150,21 +150,48 @@ func (r *Row) Xor(other *Row) *Row { } // Union returns the bitwise union of r and other. -func (r *Row) Union(other *Row) *Row { - var segments []rowSegment - itr := newMergeSegmentIterator(r.segments, other.segments) - for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() { - if s1 == nil { - segments = append(segments, *s0) - continue - } else if s0 == nil { - segments = append(segments, *s1) - continue - } - segments = append(segments, *s0.Union(s1)) +func (r *Row) Union(others ...*Row) *Row { + segments := make([][]rowSegment, 0, len(others)+1) + if len(r.segments) > 0 { + segments = append(segments, r.segments) } - - return &Row{segments: segments} + nextSegs := make([][]rowSegment, 0, len(others)+1) + toProcess := make([]*rowSegment, 0, len(others)+1) + var output []rowSegment + for _, other := range others { + if len(other.segments) > 0 { + segments = append(segments, other.segments) + } + } + for len(segments) > 0 { + shard := segments[0][0].shard + for _, segs := range segments { + if segs[0].shard < shard { + shard = segs[0].shard + } + } + nextSegs = nextSegs[:0] + toProcess := toProcess[:0] + for _, segs := range segments { + if segs[0].shard == shard { + toProcess = append(toProcess, &segs[0]) + segs = segs[1:] + } + if len(segs) > 0 { + nextSegs = append(nextSegs, segs) + } + } + // at this point, "toProcess" is a list of all the segments + // sharing the lowest ID, and nextSegs is a list of all the others. + // Swap the segment lists (so we don't have to reallocate it) + segments, nextSegs = nextSegs, segments + if len(toProcess) == 1 { + output = append(output, *toProcess[0]) + } else { + output = append(output, *toProcess[0].Union(toProcess[1:]...)) + } + } + return &Row{segments: output} } // Difference returns the diff of r and other. @@ -350,8 +377,12 @@ func (s *rowSegment) Intersect(other *rowSegment) *rowSegment { } // Union returns the bitwise union of s and other. -func (s *rowSegment) Union(other *rowSegment) *rowSegment { - data := s.data.Union(other.data) +func (s *rowSegment) Union(others ...*rowSegment) *rowSegment { + datas := make([]*roaring.Bitmap, len(others)) + for i, other := range others { + datas[i] = other.data + } + data := s.data.Union(datas...) data.Freeze() return &rowSegment{ diff --git a/server/config.go b/server/config.go index b0db09ba7..dd70adc61 100644 --- a/server/config.go +++ b/server/config.go @@ -93,6 +93,13 @@ type Config struct { // don't exhaust the goroutine limit. WorkerPoolSize int + // ImportWorkerPoolSize controls how many goroutines are created for + // processing importRoaring jobs. Defaults to runtime.NumCPU(). It is + // intentionally not defined as a flag... only exposed here so + // that we can limit the size while running tests in CI so we + // don't exhaust the goroutine limit. + ImportWorkerPoolSize int + Cluster struct { // Disabled controls whether clustering functionality is enabled. Disabled bool `toml:"disabled"` @@ -162,7 +169,8 @@ func NewConfig() *Config { TLS: TLSConfig{}, - WorkerPoolSize: runtime.NumCPU(), + WorkerPoolSize: runtime.NumCPU(), + ImportWorkerPoolSize: runtime.NumCPU(), } // Cluster config. diff --git a/server/server.go b/server/server.go index ecc2f4651..6add61591 100644 --- a/server/server.go +++ b/server/server.go @@ -317,7 +317,6 @@ func (m *Command) SetupServer() error { pilosa.OptServerMetricInterval(time.Duration(m.Config.Metric.PollInterval)), pilosa.OptServerDiagnosticsInterval(diagnosticsInterval), pilosa.OptServerExecutorPoolSize(m.Config.WorkerPoolSize), - pilosa.OptServerLogger(m.logger), pilosa.OptServerAttrStoreFunc(boltdb.NewAttrStore), pilosa.OptServerSystemInfo(gopsutil.NewSystemInfo()), @@ -347,7 +346,10 @@ func (m *Command) SetupServer() error { return errors.Wrap(err, "new server") } - m.API, err = pilosa.NewAPI(pilosa.OptAPIServer(m.Server)) + m.API, err = pilosa.NewAPI( + pilosa.OptAPIServer(m.Server), + pilosa.OptAPIImportWorkerPoolSize(m.Config.ImportWorkerPoolSize), + ) if err != nil { return errors.Wrap(err, "new api") } @@ -408,6 +410,7 @@ func (m *Command) Close() error { eg := errgroup.Group{} eg.Go(m.Handler.Close) eg.Go(m.Server.Close) + eg.Go(m.API.Close) if m.gossipMemberSet != nil { eg.Go(m.gossipMemberSet.Close) }