refactor remote mapping

build-docker.sh

@@ -1,7 +1,9 @@
-#!/bin/sh -x -e
+#!/bin/sh
+
+set -e -x
 
 GO_VER=${GO_VER:-1.5}
 
-docker run -it -v "$GOPATH":/gopath -v "$(pwd)":/app -e "GOPATH=/gopath" -w /app golang:$GO_VER sh -c 'CGO_ENABLED=0 go build -a --installsuffix cgo --ldflags="-s" -o influxd ./cmd/influxd'
+docker run -it -v "${GOPATH}":/gopath -v "$(pwd)":/app -e "GOPATH=/gopath" -w /app golang:$GO_VER sh -c 'CGO_ENABLED=0 go build -a --installsuffix cgo --ldflags="-s" -o influxd ./cmd/influxd'
 
 docker build -t influxdb .

cluster/service.go

@@ -259,6 +259,7 @@ func (s *Service) processMapShardRequest(w io.Writer, buf []byte) error {
 		if err != nil {
 			return fmt.Errorf("next chunk: %s", err)
 		}
+
 		if chunk != nil {
 			b, err := json.Marshal(chunk)
 			if err != nil {
@@ -269,11 +270,12 @@ func (s *Service) processMapShardRequest(w io.Writer, buf []byte) error {
 
 		// Write to connection.
 		resp.SetCode(0)
+		v := chunk.(*tsdb.MapperOutput)
 		if err := writeMapShardResponseMessage(w, &resp); err != nil {
 			return err
 		}
 
-		if chunk == nil {
+		if v == nil {
 			// All mapper data sent.
 			return nil
 		}

cluster/shard_mapper.go

@@ -1,7 +1,6 @@
 package cluster
 
 import (
-	"encoding/json"
 	"fmt"
 	"io"
 	"math/rand"
@@ -42,14 +41,12 @@ func NewShardMapper(timeout time.Duration) *ShardMapper {
 
 // CreateMapper returns a Mapper for the given shard ID.
 func (s *ShardMapper) CreateMapper(sh meta.ShardInfo, stmt string, chunkSize int) (tsdb.Mapper, error) {
-	var err error
-	var m tsdb.Mapper
-	if sh.OwnedBy(s.MetaStore.NodeID()) && !s.ForceRemoteMapping {
-		m, err = s.TSDBStore.CreateMapper(sh.ID, stmt, chunkSize)
-		if err != nil {
-			return nil, err
-		}
-	} else {
+	m, err := s.TSDBStore.CreateMapper(sh.ID, stmt, chunkSize)
+	if err != nil {
+		return nil, err
+	}
+
+	if !sh.OwnedBy(s.MetaStore.NodeID()) || s.ForceRemoteMapping {
 		// Pick a node in a pseudo-random manner.
 		conn, err := s.dial(sh.OwnerIDs[rand.Intn(len(sh.OwnerIDs))])
 		if err != nil {
@@ -57,8 +54,7 @@ func (s *ShardMapper) CreateMapper(sh meta.ShardInfo, stmt string, chunkSize int
 		}
 		conn.SetDeadline(time.Now().Add(s.timeout))
 
-		rm := NewRemoteMapper(conn.(*pool.PoolConn), sh.ID, stmt, chunkSize)
-		m = rm
+		m.SetRemote(NewRemoteMapper(conn.(*pool.PoolConn), sh.ID, stmt, chunkSize))
 	}
 
 	return m, nil
@@ -154,10 +150,15 @@ func (r *RemoteMapper) Open() (err error) {
 
 	// Decode the first response to get the TagSets.
 	r.tagsets = r.bufferedResponse.TagSets()
+	r.fields = r.bufferedResponse.Fields()
 
 	return nil
 }
 
+func (r *RemoteMapper) SetRemote(m tsdb.Mapper) error {
+	return fmt.Errorf("cannot set remote mapper on a remote mapper")
+}
+
 func (r *RemoteMapper) TagSets() []string {
 	return r.tagsets
 }
@@ -168,9 +169,7 @@ func (r *RemoteMapper) Fields() []string {
 
 // NextChunk returns the next chunk read from the remote node to the client.
 func (r *RemoteMapper) NextChunk() (chunk interface{}, err error) {
-	output := &tsdb.MapperOutput{}
 	var response *MapShardResponse
-
 	if r.bufferedResponse != nil {
 		response = r.bufferedResponse
 		r.bufferedResponse = nil
@@ -197,8 +196,8 @@ func (r *RemoteMapper) NextChunk() (chunk interface{}, err error) {
 	if response.Data() == nil {
 		return nil, nil
 	}
-	err = json.Unmarshal(response.Data(), output)
-	return output, err
+
+	return response.Data(), err
 }
 
 // Close the Mapper

cluster/shard_mapper_test.go

@@ -77,10 +77,14 @@ func TestShardWriter_RemoteMapper_Success(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to get next chunk from mapper: %s", err.Error())
 	}
-	output, ok := chunk.(*tsdb.MapperOutput)
+	b, ok := chunk.([]byte)
 	if !ok {
 		t.Fatal("chunk is not of expected type")
 	}
+	output := &tsdb.MapperOutput{}
+	if err := json.Unmarshal(b, output); err != nil {
+		t.Fatal(err)
+	}
 	if output.Name != "cpu" {
 		t.Fatalf("received output incorrect, exp: %v, got %v", expOutput, output)
 	}

tsdb/executor.go

@@ -24,6 +24,7 @@ const (
 // Mapper is the interface all Mapper types must implement.
 type Mapper interface {
 	Open() error
+	SetRemote(m Mapper) error
 	TagSets() []string
 	Fields() []string
 	NextChunk() (interface{}, error)

tsdb/mapper.go

@@ -3,6 +3,7 @@ package tsdb
 import (
 	"container/heap"
 	"encoding/binary"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"sort"
@@ -42,6 +43,7 @@ func (mo *MapperOutput) key() string {
 // LocalMapper is for retrieving data for a query, from a given shard.
 type LocalMapper struct {
 	shard           *Shard
+	remote          Mapper
 	stmt            influxql.Statement
 	selectStmt      *influxql.SelectStatement
 	rawMode         bool
@@ -82,6 +84,10 @@ func (lm *LocalMapper) openMeta() error {
 
 // Open opens the local mapper.
 func (lm *LocalMapper) Open() error {
+	if lm.remote != nil {
+		return lm.remote.Open()
+	}
+
 	var err error
 
 	// Get a read-only transaction.
@@ -254,10 +260,33 @@ func (lm *LocalMapper) Open() error {
 	return nil
 }
 
+func (lm *LocalMapper) SetRemote(m Mapper) error {
+	lm.remote = m
+	return nil
+}
+
 func (lm *LocalMapper) NextChunk() (interface{}, error) {
+	// If set, use remote mapper.
+	if lm.remote != nil {
+		b, err := lm.remote.NextChunk()
+		if err != nil {
+			return nil, err
+		}
+		mo := &MapperOutput{}
+		if err := json.Unmarshal(b.([]byte), mo); err != nil {
+			return nil, err
+		} else if len(mo.Values) == 0 {
+			// Mapper on other node sent 0 values so it's done.
+			return nil, nil
+		}
+		return mo, nil
+	}
+
+	// Remote mapper not set so get values from local shard.
 	if lm.rawMode {
 		return lm.nextChunkRaw()
 	}
+
 	return lm.nextChunkAgg()
 }
 
@@ -566,17 +595,27 @@ func (lm *LocalMapper) expandSources(sources influxql.Sources) (influxql.Sources
 
 // TagSets returns the list of TagSets for which this mapper has data.
 func (lm *LocalMapper) TagSets() []string {
+	if lm.remote != nil {
+		return lm.remote.TagSets()
+	}
 	return tagSetCursors(lm.cursors).Keys()
 }
 
 // Fields returns any SELECT fields. If this Mapper is not processing a SELECT query
 // then an empty slice is returned.
 func (lm *LocalMapper) Fields() []string {
+	if lm.remote != nil {
+		return lm.remote.Fields()
+	}
 	return append(lm.selectFields, lm.selectTags...)
 }
 
 // Close closes the mapper.
 func (lm *LocalMapper) Close() {
+	if lm.remote != nil {
+		lm.remote.Close()
+		return
+	}
 	if lm != nil && lm.tx != nil {
 		_ = lm.tx.Rollback()
 	}