Return results Multi and allow ordering of ops

yqwang · Dec 19, 2014 · 6b99eab · 6b99eab
1 parent 1fff4bb
commit 6b99eab
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 32 deletions.
diff --git a/zk/conn.go b/zk/conn.go
@@ -769,31 +769,40 @@ func (c *Conn) Sync(path string) (string, error) {
 	return res.Path, err
 }
 
-type MultiOps struct {
-	Create  []CreateRequest
-	Delete  []DeleteRequest
-	SetData []SetDataRequest
-	Check   []CheckVersionRequest
+type MultiResponse struct {
+	Stat   *Stat
+	String string
 }
 
-func (c *Conn) Multi(ops MultiOps) error {
+// Multi executes multiple ZooKeeper operations or none of them. The provided
+// ops must be one of *CreateRequest, *DeleteRequest, *SetDataRequest, or
+// *CheckVersionRequest.
+func (c *Conn) Multi(ops ...interface{}) ([]MultiResponse, error) {
 	req := &multiRequest{
-		Ops:        make([]multiRequestOp, 0, len(ops.Create)+len(ops.Delete)+len(ops.SetData)+len(ops.Check)),
+		Ops:        make([]multiRequestOp, 0, len(ops)),
 		DoneHeader: multiHeader{Type: -1, Done: true, Err: -1},
 	}
-	for _, r := range ops.Create {
-		req.Ops = append(req.Ops, multiRequestOp{multiHeader{opCreate, false, -1}, r})
-	}
-	for _, r := range ops.SetData {
-		req.Ops = append(req.Ops, multiRequestOp{multiHeader{opSetData, false, -1}, r})
-	}
-	for _, r := range ops.Delete {
-		req.Ops = append(req.Ops, multiRequestOp{multiHeader{opDelete, false, -1}, r})
-	}
-	for _, r := range ops.Check {
-		req.Ops = append(req.Ops, multiRequestOp{multiHeader{opCheck, false, -1}, r})
+	for _, op := range ops {
+		var opCode int32
+		switch op.(type) {
+		case *CreateRequest:
+			opCode = opCreate
+		case *SetDataRequest:
+			opCode = opSetData
+		case *DeleteRequest:
+			opCode = opDelete
+		case *CheckVersionRequest:
+			opCode = opCheck
+		default:
+			return nil, fmt.Errorf("uknown operation type %T", op)
+		}
+		req.Ops = append(req.Ops, multiRequestOp{multiHeader{opCode, false, -1}, op})
 	}
 	res := &multiResponse{}
 	_, err := c.request(opMulti, req, res, nil)
-	return err
+	mr := make([]MultiResponse, len(res.Ops))
+	for i, op := range res.Ops {
+		mr[i] = MultiResponse{Stat: op.Stat, String: op.String}
+	}
+	return mr, err
 }
diff --git a/zk/tracer.go b/zk/tracer.go
@@ -64,7 +64,6 @@ func trace(conn1, conn2 net.Conn, client bool) {
 					found := false
 					opcode, found = requests[xid]
 					if !found {
-						println("WEFWEFEW")
 						opcode = 0
 					}
 					delete(requests, xid)

diff --git a/zk/zk_test.go b/zk/zk_test.go
@@ -60,19 +60,16 @@ func TestMulti(t *testing.T) {
 	if err := zk.Delete(path, -1); err != nil && err != ErrNoNode {
 		t.Fatalf("Delete returned error: %+v", err)
 	}
-	ops := MultiOps{
-		Create: []CreateRequest{
-			{Path: path, Data: []byte{1, 2, 3, 4}, Acl: WorldACL(PermAll)},
-		},
-		SetData: []SetDataRequest{
-			{Path: path, Data: []byte{1, 2, 3, 4}, Version: -1},
-		},
-		// Delete: []DeleteRequest{
-		// 	{Path: path, Version: -1},
-		// },
-	}
-	if err := zk.Multi(ops); err != nil {
+	ops := []interface{}{
+		&CreateRequest{Path: path, Data: []byte{1, 2, 3, 4}, Acl: WorldACL(PermAll)},
+		&SetDataRequest{Path: path, Data: []byte{1, 2, 3, 4}, Version: -1},
+	}
+	if res, err := zk.Multi(ops...); err != nil {
 		t.Fatalf("Multi returned error: %+v", err)
+	} else if len(res) != 2 {
+		t.Fatalf("Expected 2 responses got %d", len(res))
+	} else {
+		t.Logf("%+v", res)
 	}
 	if data, stat, err := zk.Get(path); err != nil {
 		t.Fatalf("Get returned error: %+v", err)