Custom encoders for communication between workers via net/rpc over TCP. Also, hello there!

conn/client.go

+package conn
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"net/rpc"
+)
+
+type ClientCodec struct {
+	Rwc        io.ReadWriteCloser
+	payloadLen int32
+}
+
+func (c *ClientCodec) WriteRequest(r *rpc.Request, body interface{}) error {
+	if body == nil {
+		return fmt.Errorf("Nil request body from client.")
+	}
+
+	query := body.(*Query)
+	if err := writeHeader(c.Rwc, r.Seq, r.ServiceMethod, query.Data); err != nil {
+		return err
+	}
+	n, err := c.Rwc.Write(query.Data)
+	if n != len(query.Data) {
+		return errors.New("Unable to write payload.")
+	}
+	return err
+}
+
+func (c *ClientCodec) ReadResponseHeader(r *rpc.Response) error {
+	if len(r.Error) > 0 {
+		log.Fatal("client got response error: " + r.Error)
+	}
+	if err := parseHeader(c.Rwc, &r.Seq,
+		&r.ServiceMethod, &c.payloadLen); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (c *ClientCodec) ReadResponseBody(body interface{}) error {
+	buf := make([]byte, c.payloadLen)
+	n, err := c.Rwc.Read(buf)
+	if n != int(c.payloadLen) {
+		return fmt.Errorf("ClientCodec expected: %d. Got: %d\n", c.payloadLen, n)
+	}
+
+	reply := body.(*Reply)
+	reply.Data = buf
+	return err
+}
+
+func (c *ClientCodec) Close() error {
+	return c.Rwc.Close()
+}

conn/codec.go

+package conn
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"io"
+
+	"github.com/dgraph-io/dgraph/x"
+)
+
+type Query struct {
+	Data []byte
+}
+
+type Reply struct {
+	Data []byte
+}
+
+func writeHeader(rwc io.ReadWriteCloser, seq uint64,
+	method string, data []byte) error {
+
+	var bh bytes.Buffer
+	var rerr error
+
+	x.SetError(&rerr, binary.Write(&bh, binary.LittleEndian, seq))
+	x.SetError(&rerr, binary.Write(&bh, binary.LittleEndian, int32(len(method))))
+	x.SetError(&rerr, binary.Write(&bh, binary.LittleEndian, int32(len(data))))
+	_, err := bh.Write([]byte(method))
+	x.SetError(&rerr, err)
+	if rerr != nil {
+		return rerr
+	}
+	_, err = rwc.Write(bh.Bytes())
+	return err
+}
+
+func parseHeader(rwc io.ReadWriteCloser, seq *uint64,
+	method *string, plen *int32) error {
+
+	var err error
+	var sz int32
+	x.SetError(&err, binary.Read(rwc, binary.LittleEndian, seq))
+	x.SetError(&err, binary.Read(rwc, binary.LittleEndian, &sz))
+	x.SetError(&err, binary.Read(rwc, binary.LittleEndian, plen))
+	if err != nil {
+		return err
+	}
+
+	buf := make([]byte, sz)
+	n, err := rwc.Read(buf)
+	if err != nil {
+		return err
+	}
+	if n != int(sz) {
+		return fmt.Errorf("Expected: %v. Got: %v\n", sz, n)
+	}
+	*method = string(buf)
+	return nil
+}

conn/pool.go

+package conn
+
+import (
+	"net"
+	"net/rpc"
+
+	"github.com/dgraph-io/dgraph/x"
+)
+
+var glog = x.Log("conn")
+
+type Pool struct {
+	clients chan *rpc.Client
+	addr    string
+}
+
+func NewPool(addr string, maxCap int) *Pool {
+	p := new(Pool)
+	p.addr = addr
+	p.clients = make(chan *rpc.Client, maxCap)
+	client, err := p.dialNew()
+	if err != nil {
+		glog.Fatal(err)
+		return nil
+	}
+	p.clients <- client
+	return p
+}
+
+func (p *Pool) dialNew() (*rpc.Client, error) {
+	nconn, err := net.Dial("tcp", p.addr)
+	if err != nil {
+		return nil, err
+	}
+	cc := &ClientCodec{
+		Rwc: nconn,
+	}
+	return rpc.NewClientWithCodec(cc), nil
+}
+
+func (p *Pool) Get() (*rpc.Client, error) {
+	select {
+	case client := <-p.clients:
+		return client, nil
+	default:
+		return p.dialNew()
+	}
+}
+
+func (p *Pool) Put(client *rpc.Client) error {
+	select {
+	case p.clients <- client:
+		return nil
+	default:
+		return client.Close()
+	}
+}
+
+func (p *Pool) Close() error {
+	// We're not doing a clean exit here. A clean exit here would require
+	// mutex locks around conns; which seems unnecessary just to shut down
+	// the server.
+	return nil
+}

test/server.go → conn/server.go

-package main
+package conn
 
 import (
-	"bufio"
 	"errors"
-	"fmt"
 	"io"
 	"log"
 	"net/rpc"
-	"reflect"
 )
 
-type scodec struct {
-	rwc        io.ReadWriteCloser
-	ebuf       *bufio.Writer
+type ServerCodec struct {
+	Rwc        io.ReadWriteCloser
 	payloadLen int32
 }
 
-func (c *scodec) ReadRequestHeader(r *rpc.Request) error {
-	var err error
-	if err = parseHeader(c.rwc, &r.Seq,
-		&r.ServiceMethod, &c.payloadLen); err != nil {
-		return err
-	}
-
-	fmt.Println("server using custom codec to read header")
-	fmt.Println("server method called:", r.ServiceMethod)
-	fmt.Println("server method called:", r.Seq)
-	return nil
+func (c *ServerCodec) ReadRequestHeader(r *rpc.Request) error {
+	return parseHeader(c.Rwc, &r.Seq, &r.ServiceMethod, &c.payloadLen)
 }
 
-func (c *scodec) ReadRequestBody(data interface{}) error {
-	if data == nil {
-		log.Fatal("Why is data nil here?")
-	}
-	value := reflect.ValueOf(data)
-	if value.Type().Kind() != reflect.Ptr {
-		log.Fatal("Should of of type pointer")
-	}
-
+func (c *ServerCodec) ReadRequestBody(data interface{}) error {
 	b := make([]byte, c.payloadLen)
-	n, err := c.rwc.Read(b)
-	fmt.Printf("Worker read n bytes: %v %s\n", n, string(b))
+	n, err := c.Rwc.Read(b)
 	if err != nil {
 		log.Fatal("server", err)
 	}
 	if n != int(c.payloadLen) {
-		return errors.New("Server unable to read request.")
+		return errors.New("ServerCodec unable to read request.")
 	}
 
+	if data == nil {
+		// If data is nil, discard this request.
+		return nil
+	}
 	query := data.(*Query)
-	query.d = b
+	query.Data = b
 	return nil
 }
 
-func (c *scodec) WriteResponse(resp *rpc.Response, data interface{}) error {
+func (c *ServerCodec) WriteResponse(resp *rpc.Response, data interface{}) error {
 	if len(resp.Error) > 0 {
 		log.Fatal("Response has error: " + resp.Error)
 	}
@@ -65,15 +47,15 @@ func (c *scodec) WriteResponse(resp *rpc.Response, data interface{}) error {
 		log.Fatal("Unable to convert to reply")
 	}
 
-	if err := writeHeader(c.rwc, resp.Seq,
-		resp.ServiceMethod, reply.d); err != nil {
+	if err := writeHeader(c.Rwc, resp.Seq,
+		resp.ServiceMethod, reply.Data); err != nil {
 		return err
 	}
 
-	_, err := c.rwc.Write(reply.d)
+	_, err := c.Rwc.Write(reply.Data)
 	return err
 }
 
-func (c *scodec) Close() error {
-	return c.rwc.Close()
+func (c *ServerCodec) Close() error {
+	return c.Rwc.Close()
 }

server/main.go

@@ -141,6 +141,7 @@ func main() {
 
 	posting.Init(clog)
 	worker.Init(ps)
+	worker.Connect()
 	uid.Init(ps)
 
 	http.HandleFunc("/query", queryHandler(ps))

task.fbs

@@ -9,6 +9,10 @@ table Value {
 	val:[ubyte];
 }
 
+table XidList {
+	xids:[string];
+}
+
 table UidList {
 	uids:[ulong];
 }

task/XidList.go

+// automatically generated, do not modify
+
+package task
+
+import (
+	flatbuffers "github.com/google/flatbuffers/go"
+)
+type XidList struct {
+	_tab flatbuffers.Table
+}
+
+func (rcv *XidList) Init(buf []byte, i flatbuffers.UOffsetT) {
+	rcv._tab.Bytes = buf
+	rcv._tab.Pos = i
+}
+
+func (rcv *XidList) Xids(j int) []byte {
+	o := flatbuffers.UOffsetT(rcv._tab.Offset(4))
+	if o != 0 {
+		a := rcv._tab.Vector(o)
+		return rcv._tab.ByteVector(a + flatbuffers.UOffsetT(j * 4))
+	}
+	return nil
+}
+
+func (rcv *XidList) XidsLength() int {
+	o := flatbuffers.UOffsetT(rcv._tab.Offset(4))
+	if o != 0 {
+		return rcv._tab.VectorLen(o)
+	}
+	return 0
+}
+
+func XidListStart(builder *flatbuffers.Builder) { builder.StartObject(1) }
+func XidListAddXids(builder *flatbuffers.Builder, xids flatbuffers.UOffsetT) { builder.PrependUOffsetTSlot(0, flatbuffers.UOffsetT(xids), 0) }
+func XidListStartXidsVector(builder *flatbuffers.Builder, numElems int) flatbuffers.UOffsetT { return builder.StartVector(4, numElems, 4)
+}
+func XidListEnd(builder *flatbuffers.Builder) flatbuffers.UOffsetT { return builder.EndObject() }

test/client.go

-package main
-
-import (
-	"bufio"
-	"bytes"
-	"encoding/binary"
-	"errors"
-	"fmt"
-	"io"
-	"log"
-	"net/rpc"
-)
-
-type ccodec struct {
-	rwc        io.ReadWriteCloser
-	ebuf       *bufio.Writer
-	payloadLen int32
-}
-
-func writeHeader(rwc io.ReadWriteCloser, seq uint64,
-	method string, data []byte) error {
-
-	var bh bytes.Buffer
-	var rerr error
-
-	setError(&rerr, binary.Write(&bh, binary.LittleEndian, seq))
-	setError(&rerr, binary.Write(&bh, binary.LittleEndian, int32(len(method))))
-	setError(&rerr, binary.Write(&bh, binary.LittleEndian, int32(len(data))))
-	_, err := bh.Write([]byte(method))
-	setError(&rerr, err)
-	if rerr != nil {
-		return rerr
-	}
-	_, err = rwc.Write(bh.Bytes())
-	return err
-}
-
-func parseHeader(rwc io.ReadWriteCloser, seq *uint64, method *string, plen *int32) error {
-	var err error
-	var sz int32
-	setError(&err, binary.Read(rwc, binary.LittleEndian, seq))
-	setError(&err, binary.Read(rwc, binary.LittleEndian, &sz))
-	setError(&err, binary.Read(rwc, binary.LittleEndian, plen))
-	if err != nil {
-		return err
-	}
-	buf := make([]byte, sz)
-	n, err := rwc.Read(buf)
-	if err != nil {
-		return err
-	}
-	if n != int(sz) {
-		return fmt.Errorf("Expected: %v. Got: %v\n", sz, n)
-	}
-	*method = string(buf)
-	return nil
-}
-
-func (c *ccodec) WriteRequest(r *rpc.Request, body interface{}) error {
-	if body == nil {
-		return errors.New("Nil body")
-	}
-
-	query := body.(*Query)
-	if err := writeHeader(c.rwc, r.Seq, r.ServiceMethod, query.d); err != nil {
-		return err
-	}
-
-	n, err := c.rwc.Write(query.d)
-	if n != len(query.d) {
-		return errors.New("Unable to write payload.")
-	}
-	return err
-}
-
-func (c *ccodec) ReadResponseHeader(r *rpc.Response) error {
-	if len(r.Error) > 0 {
-		log.Fatal("client got response error: " + r.Error)
-	}
-	if err := parseHeader(c.rwc, &r.Seq,
-		&r.ServiceMethod, &c.payloadLen); err != nil {
-		return err
-	}
-	fmt.Println("Client got response:", r.Seq)
-	fmt.Println("Client got response:", r.ServiceMethod)
-	return nil
-}
-
-func (c *ccodec) ReadResponseBody(body interface{}) error {
-	buf := make([]byte, c.payloadLen)
-	n, err := c.rwc.Read(buf)
-	if n != int(c.payloadLen) {
-		return fmt.Errorf("Client expected: %d. Got: %d\n", c.payloadLen, n)
-	}
-	reply := body.(*Reply)
-	reply.d = buf
-	return err
-}
-
-func (c *ccodec) Close() error {
-	return c.rwc.Close()
-}

test/main.go

-package main
-
-import (
-	"bufio"
-	"fmt"
-	"io"
-	"log"
-	"math/rand"
-	"net"
-	"net/rpc"
-)
-
-type Query struct {
-	d []byte
-}
-
-type Reply struct {
-	d []byte
-}
-
-func setError(prev *error, n error) {
-	if prev == nil {
-		prev = &n
-	}
-}
-
-type Worker struct {
-}
-
-func serveIt(conn io.ReadWriteCloser) {
-	for {
-		srv := &scodec{
-			rwc:  conn,
-			ebuf: bufio.NewWriter(conn),
-		}
-		rpc.ServeRequest(srv)
-	}
-}
-
-func (w *Worker) Receive(query *Query, reply *Reply) error {
-	fmt.Printf("Worker received: [%s]\n", string(query.d))
-	reply.d = []byte("abcdefghij-Hello World!")
-	return nil
-}
-
-func runServer(address string) error {
-	w := new(Worker)
-	if err := rpc.Register(w); err != nil {
-		return err
-	}
-
-	ln, err := net.Listen("tcp", address)
-	if err != nil {
-		fmt.Printf("listen(%q): %s\n", address, err)
-		return err
-	}
-	fmt.Printf("Worker listening on %s\n", ln.Addr())
-	go func() {
-		for {
-			cxn, err := ln.Accept()
-			if err != nil {
-				log.Fatalf("listen(%q): %s\n", address, err)
-				return
-			}
-			log.Printf("Worker accepted connection to %s from %s\n",
-				cxn.LocalAddr(), cxn.RemoteAddr())
-			go serveIt(cxn)
-		}
-	}()
-	return nil
-}
-
-func main() {
-	addresses := map[int]string{
-		1: "127.0.0.1:10000",
-		// 2: "127.0.0.1:10001",
-		// 3: "127.0.0.1:10002",
-	}
-
-	for _, address := range addresses {
-		runServer(address)
-	}
-
-	clients := make(map[int]*rpc.Client)
-	for id, address := range addresses {
-		conn, err := net.Dial("tcp", address)
-		if err != nil {
-			log.Fatal("dial", err)
-		}
-		cc := &ccodec{
-			rwc:  conn,
-			ebuf: bufio.NewWriter(conn),
-		}
-		clients[id] = rpc.NewClientWithCodec(cc)
-	}
-
-	for i := 0; i < 1; i++ {
-		client := clients[1]
-		if client == nil {
-			log.Fatal("Worker is nil")
-		}
-
-		id := 0
-		// for id, server := range servers {
-		query := new(Query)
-		query.d = []byte(fmt.Sprintf("id:%d Rand: %d", id, rand.Int()))
-		reply := new(Reply)
-		if err := client.Call("Worker.Receive", query, reply); err != nil {
-			log.Fatal("call", err)
-		}
-
-		fmt.Printf("Returned: %s\n", string(reply.d))
-		// }
-	}
-}

worker/worker.go

 package worker
 
 import (
+	"flag"
+	"io"
+	"net"
+	"net/rpc"
+	"strings"
+
+	"github.com/dgraph-io/dgraph/conn"
 	"github.com/dgraph-io/dgraph/posting"
 	"github.com/dgraph-io/dgraph/store"
 	"github.com/dgraph-io/dgraph/task"
@@ -8,12 +15,53 @@ import (
 	"github.com/google/flatbuffers/go"
 )
 
+var workers = flag.String("workers", "",
+	"Comma separated list of IP addresses of workers")
+var workerPort = flag.String("workerport", ":12345",
+	"Port used by worker for internal communication.")
+
+var glog = x.Log("worker")
 var dataStore *store.Store
+var pools []*conn.Pool
 
 func Init(ps *store.Store) {
 	dataStore = ps
 }
 
+func Connect() {
+	w := new(Worker)
+	if err := rpc.Register(w); err != nil {
+		glog.Fatal(err)
+	}
+	if err := runServer(*workerPort); err != nil {
+		glog.Fatal(err)
+	}
+
+	addrs := strings.Split(*workers, ",")
+	var pools []*conn.Pool
+	for _, addr := range addrs {
+		if len(addr) == 0 {
+			continue
+		}
+		pool := conn.NewPool(addr, 5)
+		client, err := pool.Get()
+		if err != nil {
+			glog.Fatal(err)
+		}
+		query := new(conn.Query)
+		query.Data = []byte("hello")
+		reply := new(conn.Reply)
+		if err = client.Call("Worker.Hello", query, reply); err != nil {
+			glog.WithField("call", "Worker.Hello").Fatal(err)
+		}
+		glog.WithField("reply", string(reply.Data)).WithField("addr", addr).
+			Info("Got reply from server")
+		pools = append(pools, pool)
+	}
+
+	glog.Info("Server started. Clients connected.")
+}
+
 func ProcessTask(query []byte) (result []byte, rerr error) {
 	uo := flatbuffers.GetUOffsetT(query)
 	q := new(task.Query)
@@ -78,3 +126,48 @@ func NewQuery(attr string, uids []uint64) []byte {
 	b.Finish(qend)
 	return b.Bytes[b.Head():]
 }
+
+type Worker struct {
+}
+
+func (w *Worker) Hello(query *conn.Query, reply *conn.Reply) error {
+	if string(query.Data) == "hello" {
+		reply.Data = []byte("Oh hello there!")
+	} else {
+		reply.Data = []byte("Hey stranger!")
+	}
+	return nil
+}
+
+func serveRequests(irwc io.ReadWriteCloser) {
+	for {
+		sc := &conn.ServerCodec{
+			Rwc: irwc,
+		}
+		rpc.ServeRequest(sc)
+	}
+}
+
+func runServer(address string) error {
+	ln, err := net.Listen("tcp", address)
+	if err != nil {
+		glog.Fatalf("While running server: %v", err)
+		return err
+	}
+	glog.WithField("address", ln.Addr()).Info("Worker listening")
+
+	go func() {
+		for {
+			cxn, err := ln.Accept()
+			if err != nil {
+				glog.Fatalf("listen(%q): %s\n", address, err)
+				return
+			}
+			glog.WithField("local", cxn.LocalAddr()).
+				WithField("remote", cxn.RemoteAddr()).
+				Debug("Worker accepted connection")
+			go serveRequests(cxn)
+		}
+	}()
+	return nil
+}

x/x.go

@@ -39,6 +39,12 @@ type DirectedEdge struct {
 	Timestamp time.Time
 }
 
+func SetError(prev *error, n error) {
+	if prev == nil {
+		prev = &n
+	}
+}
+
 func Log(p string) *logrus.Entry {
 	l := logrus.WithFields(logrus.Fields{
 		"package": p,