package pilosa import ( "bytes" "errors" "fmt" "io/ioutil" "net/http" "net/url" "sort" "time" "github.com/gogo/protobuf/proto" "github.com/umbel/pilosa/internal" "github.com/umbel/pilosa/pql" ) // DefaultFrame is the frame used if one is not specified. const DefaultFrame = "general" // Executor recursively executes calls in a PQL query across all slices. type Executor struct { index *Index // Local hostname & cluster configuration. Host string Cluster *Cluster // Client used for remote HTTP requests. HTTPClient *http.Client } // NewExecutor returns a new instance of Executor. func NewExecutor(index *Index) *Executor { return &Executor{ index: index, HTTPClient: http.DefaultClient, } } // Index returns the index that the executor runs against. func (e *Executor) Index() *Index { return e.index } // Execute executes a PQL query. func (e *Executor) Execute(db string, q *pql.Query, slices []uint64, opt *ExecOptions) (interface{}, error) { // Verify that a database is set. if db == "" { return nil, ErrDatabaseRequired } // Default options. if opt == nil { opt = &ExecOptions{} } // Ignore slices for set calls. switch root := q.Root.(type) { case *pql.SetBit: return e.executeSetBit(db, root, opt) case *pql.SetBitmapAttrs: return nil, e.executeSetBitmapAttrs(db, root) case *pql.SetProfileAttrs: return nil, e.executeSetProfileAttrs(db, root) } // If slices aren't specified, then include all of them. if len(slices) == 0 { // Round up the number of slices. sliceN := e.index.SliceN() sliceN += (sliceN % uint64(len(e.Cluster.Nodes))) + uint64(len(e.Cluster.Nodes)) // Generate a slices of all slices. slices = make([]uint64, sliceN+1) for i := range slices { slices[i] = uint64(i) } } return e.executeCall(db, q.Root, slices, opt) } // executeCall executes a call. func (e *Executor) executeCall(db string, c pql.Call, slices []uint64, opt *ExecOptions) (interface{}, error) { switch c := c.(type) { case pql.BitmapCall: return e.executeBitmapCall(db, c, slices, opt) case *pql.Count: return e.executeCount(db, c, slices, opt) case *pql.Profile: return e.executeProfile(db, c, opt) case *pql.TopN: return e.executeTopN(db, c, slices, opt) default: panic("unreachable") } } // executeBitmapCall executes a call that returns a bitmap. func (e *Executor) executeBitmapCall(db string, c pql.BitmapCall, slices []uint64, opt *ExecOptions) (*Bitmap, error) { other := NewBitmap() for node, nodeSlices := range e.slicesByNode(slices) { // Execute locally if the hostname matches. if node.Host == e.Host { for _, slice := range nodeSlices { bm, err := e.executeBitmapCallSlice(db, c, slice) if err != nil { return nil, err } other.Merge(bm) } continue } // Otherwise execute remotely. res, err := e.exec(node, db, &pql.Query{Root: c}, nodeSlices, opt) if err != nil { return nil, err } other.Merge(res.(*Bitmap)) } // Attach bitmap attributes for Bitmap() calls. if c, ok := c.(*pql.Bitmap); ok { fr := e.Index().Frame(db, c.Frame) if fr != nil { attrs, err := fr.BitmapAttrStore().Attrs(c.ID) if err != nil { return nil, err } other.Attrs = attrs } } return other, nil } // executeBitmapCallSlice executes a bitmap call for a single slice. func (e *Executor) executeBitmapCallSlice(db string, c pql.BitmapCall, slice uint64) (*Bitmap, error) { switch c := c.(type) { case *pql.Bitmap: return e.executeBitmapSlice(db, c, slice) case *pql.Difference: return e.executeDifferenceSlice(db, c, slice) case *pql.Intersect: return e.executeIntersectSlice(db, c, slice) case *pql.Range: return e.executeRangeSlice(db, c, slice) case *pql.Union: return e.executeUnionSlice(db, c, slice) default: panic("unreachable") } } // executeTopN executes a TopN() call. func (e *Executor) executeTopN(db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) { var results []Pair for node, nodeSlices := range e.slicesByNode(slices) { // Execute locally if the hostname matches. if node.Host == e.Host { for _, slice := range nodeSlices { pairs, err := e.executeTopNSlice(db, c, slice) if err != nil { return nil, err } results = Pairs(results).Add(pairs) } continue } // Otherwise execute remotely. res, err := e.exec(node, db, &pql.Query{Root: c}, nodeSlices, opt) if err != nil { return nil, err } results = Pairs(results).Add(res.([]Pair)) } // Sort final merged results. sort.Sort(Pairs(results)) // Only keep the top n after sorting. if len(results) > c.N { results = results[0:c.N] } return results, nil } // executeTopNSlice executes a TopN call for a single slice. func (e *Executor) executeTopNSlice(db string, c *pql.TopN, slice uint64) ([]Pair, error) { // Retrieve bitmap used to intersect. var src *Bitmap if c.Src != nil { bm, err := e.executeBitmapCallSlice(db, c.Src, slice) if err != nil { return nil, err } src = bm } // Set default frame. frame := c.Frame if frame == "" { frame = DefaultFrame } f := e.Index().Fragment(db, frame, slice) if f == nil { return nil, nil } return f.TopN(c.N, src, c.Field, c.Filters) } // executeDifferenceSlice executes a difference() call for a local slice. func (e *Executor) executeDifferenceSlice(db string, c *pql.Difference, slice uint64) (*Bitmap, error) { var other *Bitmap for i, input := range c.Inputs { bm, err := e.executeBitmapCallSlice(db, input, slice) if err != nil { return nil, err } if i == 0 { other = bm } else { other = other.Difference(bm) } } other.SetCount(other.BitCount()) return other, nil } func (e *Executor) executeBitmapSlice(db string, c *pql.Bitmap, slice uint64) (*Bitmap, error) { frame := c.Frame if frame == "" { frame = DefaultFrame } f := e.Index().Fragment(db, frame, slice) if f == nil { return NewBitmap(), nil } return f.Bitmap(c.ID), nil } // executeIntersectSlice executes a intersect() call for a local slice. func (e *Executor) executeIntersectSlice(db string, c *pql.Intersect, slice uint64) (*Bitmap, error) { var other *Bitmap for i, input := range c.Inputs { bm, err := e.executeBitmapCallSlice(db, input, slice) if err != nil { return nil, err } if i == 0 { other = bm } else { other = other.Intersect(bm) } } other.SetCount(other.BitCount()) return other, nil } // executeRangeSlice executes a range() call for a local slice. func (e *Executor) executeRangeSlice(db string, c *pql.Range, slice uint64) (*Bitmap, error) { frame := c.Frame if frame == "" { frame = DefaultFrame } f := e.Index().Fragment(db, frame, slice) if f == nil { return NewBitmap(), nil } return f.Range(c.ID, c.StartTime, c.EndTime), nil } // executeUnionSlice executes a union() call for a local slice. func (e *Executor) executeUnionSlice(db string, c *pql.Union, slice uint64) (*Bitmap, error) { var other *Bitmap for i, input := range c.Inputs { bm, err := e.executeBitmapCallSlice(db, input, slice) if err != nil { return nil, err } if i == 0 { other = bm } else { other = other.Union(bm) } } other.SetCount(other.BitCount()) return other, nil } // executeCount executes a count() call. func (e *Executor) executeCount(db string, c *pql.Count, slices []uint64, opt *ExecOptions) (uint64, error) { var n uint64 for node, nodeSlices := range e.slicesByNode(slices) { // Execute locally if the hostname matches. if node.Host == e.Host { for _, slice := range nodeSlices { bm, err := e.executeBitmapCallSlice(db, c.Input, slice) if err != nil { return 0, err } n += bm.Count() } continue } // Otherwise execute remotely. res, err := e.exec(node, db, &pql.Query{Root: c}, nodeSlices, opt) if err != nil { return 0, err } n += res.(uint64) } return n, nil } // executeProfile executes a Profile() call. // This call only executes locally since the profile attibutes are stored locally. func (e *Executor) executeProfile(db string, c *pql.Profile, opt *ExecOptions) (*Profile, error) { panic("FIXME: impl: e.Index().ProfileAttr(c.ID)") } // executeSetBit executes a SetBit() call. func (e *Executor) executeSetBit(db string, c *pql.SetBit, opt *ExecOptions) (bool, error) { slice := c.ProfileID / SliceWidth ret := false for _, node := range e.Cluster.SliceNodes(slice) { // Update locally if host matches. if node.Host == e.Host { f, err := e.Index().CreateFragmentIfNotExists(db, c.Frame, slice) if err != nil { return false, fmt.Errorf("fragment: %s", err) } val, err := f.SetBit(c.ID, c.ProfileID, opt.Timestamp, opt.Quantum) if err != nil { return false, err } if val { ret = true } continue } // Forward call to remote node otherwise. if res, err := e.exec(node, db, &pql.Query{Root: c}, nil, opt); err != nil { return false, err } else { ret = res.(bool) } } return ret, nil } // executeSetBitmapAttrs executes a SetBitmapAttrs() call. func (e *Executor) executeSetBitmapAttrs(db string, c *pql.SetBitmapAttrs) error { // Retrieve frame. frame, err := e.Index().CreateFrameIfNotExists(db, c.Frame) if err != nil { return err } // Set attributes. if err := frame.BitmapAttrStore().SetAttrs(c.ID, c.Attrs); err != nil { return err } // TODO: Propagate attributes to other servers in cluster. return nil } // executeSetProfileAttrs executes a SetProfileAttrs() call. func (e *Executor) executeSetProfileAttrs(db string, c *pql.SetProfileAttrs) error { // Retrieve database. d, err := e.Index().CreateDBIfNotExists(db) if err != nil { return err } // Set attributes. if err := d.ProfileAttrStore().SetAttrs(c.ID, c.Attrs); err != nil { return err } // TODO: Propagate attributes to other servers in cluster. return nil } // exec executes a PQL query remotely for a set of slices on a node. func (e *Executor) exec(node *Node, db string, q *pql.Query, slices []uint64, opt *ExecOptions) (result interface{}, err error) { // Encode request object. pbreq := &internal.QueryRequest{ DB: proto.String(db), Query: proto.String(q.String()), Slices: slices, Quantum: proto.Uint32(uint32(opt.Quantum)), } if opt.Timestamp != nil { pbreq.Timestamp = proto.Int64(opt.Timestamp.UnixNano()) } buf, err := proto.Marshal(pbreq) if err != nil { return nil, err } // Create HTTP request. req, err := http.NewRequest("POST", (&url.URL{ Scheme: "http", Host: node.Host, Path: "/query", }).String(), bytes.NewReader(buf)) if err != nil { return nil, err } // Require protobuf encoding. req.Header.Set("Accept", "application/x-protobuf") req.Header.Set("Content-Type", "application/x-protobuf") // Send request to remote node. resp, err := e.HTTPClient.Do(req) if err != nil { return nil, err } defer resp.Body.Close() // Read response into buffer. body, err := ioutil.ReadAll(resp.Body) if err != nil { return nil, err } // Check status code. if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("invalid status: code=%d, err=%s", resp.StatusCode, body) } // Decode response object. var pb internal.QueryResponse if err := proto.Unmarshal(body, &pb); err != nil { return nil, err } // Return an error, if specified on response. if err := decodeError(pb.GetErr()); err != nil { return nil, err } // Return appropriate data for the query. switch q.Root.(type) { case pql.BitmapCall: return decodeBitmap(pb.GetBitmap()), nil case *pql.TopN: return decodePairs(pb.GetPairs()), nil case *pql.Count: return pb.GetN(), nil case *pql.SetBit: return pb.GetChanged(), nil default: panic(fmt.Sprintf("invalid node for remote exec: %T", q.Root)) } } // slicesByNode returns a mapping of nodes to slices. // // NOTE: Currently the only primary node is used. func (e *Executor) slicesByNode(slices []uint64) map[*Node][]uint64 { m := make(map[*Node][]uint64) for _, slice := range slices { nodes := e.Cluster.SliceNodes(slice) node := nodes[0] m[node] = append(m[node], slice) } return m } // ExecOptions represents an execution context for a single Execute() call. type ExecOptions struct { Timestamp *time.Time Quantum TimeQuantum } // decodeError returns an error representation of s if s is non-blank. // Returns nil if s is blank. func decodeError(s string) error { if s == "" { return nil } return errors.New(s) }