rpc: allow served object to know about shutdown.

rpc/rpc_test.go

 package rpc_test
 
 import (
         "encoding/json"
         "fmt"
         "io"
         . "launchpad.net/gocheck"
         "launchpad.net/juju-core/log"
         "launchpad.net/juju-core/rpc"
         "launchpad.net/juju-core/testing"
@@ skipping 23 matching lines @@
 type callError callInfo
 
 func (e *callError) Error() string {
         return fmt.Sprintf("error calling %s", e.method)
 }
 
 type stringVal struct {
         Val string
 }
 
-type TRoot struct {
+type Root struct {
         mu        sync.Mutex
         calls     []*callInfo
         returnErr bool
         simple    map[string]*SimpleMethods
         delayed   map[string]*DelayedMethods
         errorInst *ErrorMethods
 }
 
-func (r *TRoot) callError(rcvr interface{}, name string, arg interface{}) error {
+func (r *Root) callError(rcvr interface{}, name string, arg interface{}) error {
         if r.returnErr {
                 return &callError{rcvr, name, arg}
         }
         return nil
 }
 
-func (r *TRoot) SimpleMethods(id string) (*SimpleMethods, error) {
+func (r *Root) SimpleMethods(id string) (*SimpleMethods, error) {
         r.mu.Lock()
         defer r.mu.Unlock()
         if a := r.simple[id]; a != nil {
                 return a, nil
         }
         return nil, fmt.Errorf("unknown SimpleMethods id")
 }
 
-func (r *TRoot) DelayedMethods(id string) (*DelayedMethods, error) {
+func (r *Root) DelayedMethods(id string) (*DelayedMethods, error) {
         r.mu.Lock()
         defer r.mu.Unlock()
         if a := r.delayed[id]; a != nil {
                 return a, nil
         }
         return nil, fmt.Errorf("unknown DelayedMethods id")
 }
 
-func (r *TRoot) ErrorMethods(id string) (*ErrorMethods, error) {
+func (r *Root) ErrorMethods(id string) (*ErrorMethods, error) {
         if r.errorInst == nil {
                 return nil, fmt.Errorf("no error methods")
         }
         return r.errorInst, nil
 }
 
-func (t *TRoot) called(rcvr interface{}, method string, arg interface{}) {
+func (t *Root) called(rcvr interface{}, method string, arg interface{}) {
         t.mu.Lock()
         t.calls = append(t.calls, &callInfo{rcvr, method, arg})
         t.mu.Unlock()
 }
 
 type SimpleMethods struct {
-        root *TRoot
+        root *Root
         id   string
 }
 
 // Each Call method is named in this standard form:
 //
 //     Call<narg>r<nret><e>
 //
 // where narg is the number of arguments, nret is the number of returned
 // values (not including the error) and e is the letter 'e' if the
 // method returns an error.
@@ skipping 50 matching lines @@
 
 type ErrorMethods struct {
         err error
 }
 
 func (e *ErrorMethods) Call() error {
         return e.err
 }
 
 func (*suite) TestRPC(c *C) {
-        root := &TRoot{
+        root := &Root{
                 simple: make(map[string]*SimpleMethods),
         }
         root.simple["a99"] = &SimpleMethods{root: root, id: "a99"}
         client, srvDone := newRPCClientServer(c, root, nil)
         for narg := 0; narg < 2; narg++ {
                 for nret := 0; nret < 2; nret++ {
                         for nerr := 0; nerr < 2; nerr++ {
                                 retErr := nerr != 0
                                 root.testCall(c, client, narg, nret, retErr, false)
                                 if retErr {
                                         root.testCall(c, client, narg, nret, retErr, true)
                                 }
                         }
                 }
         }
         client.Close()
         err := chanReadError(c, srvDone, "server done")
         c.Assert(err, IsNil)
 }
 
-func (root *TRoot) testCall(c *C, client *rpc.Client, narg, nret int, retErr, testErr bool) {
+func (root *Root) testCall(c *C, client *rpc.Client, narg, nret int, retErr, testErr bool) {
         root.calls = nil
         root.returnErr = testErr
         e := ""
         if retErr {
                 e = "e"
         }
         method := fmt.Sprintf("Call%dr%d%s", narg, nret, e)
         c.Logf("test call %s", method)
         var r stringVal
         err := client.Call("SimpleMethods", "a99", method, stringVal{"arg"}, &r)
@@ skipping 18 matching lines @@
                 c.Assert(r, Equals, stringVal{method + " ret"})
         }
 }
 
 func (*suite) TestConcurrentCalls(c *C) {
         start1 := make(chan string)
         start2 := make(chan string)
         ready1 := make(chan struct{})
         ready2 := make(chan struct{})
 
-        root := &TRoot{
+        root := &Root{
                 delayed: map[string]*DelayedMethods{
                         "1": {ready: ready1, done: start1},
                         "2": {ready: ready2, done: start2},
                 },
         }
 
         client, srvDone := newRPCClientServer(c, root, nil)
         call := func(id string, done chan<- struct{}) {
                 var r stringVal
                 err := client.Call("DelayedMethods", id, "Delay", nil, &r)
@@ skipping 27 matching lines @@
 
 func (e *codedError) Error() string {
         return e.m
 }
 
 func (e *codedError) ErrorCode() string {
         return e.code
 }
 
 func (*suite) TestErrorCode(c *C) {
-        root := &TRoot{
+        root := &Root{
                 errorInst: &ErrorMethods{&codedError{"message", "code"}},
         }
         client, srvDone := newRPCClientServer(c, root, nil)
         err := client.Call("ErrorMethods", "", "Call", nil, nil)
         c.Assert(err, ErrorMatches, `server error: message \(code\)`)
         c.Assert(err.(rpc.ErrorCoder).ErrorCode(), Equals, "code")
         client.Close()
         err = chanReadError(c, srvDone, "server done")
         c.Assert(err, IsNil)
 }
 
 func (*suite) TestTransformErrors(c *C) {
-        root := &TRoot{
+        root := &Root{
                 errorInst: &ErrorMethods{&codedError{"message", "code"}},
         }
         tfErr := func(err error) error {
                 c.Check(err, NotNil)
                 if e, ok := err.(*codedError); ok {
                         return &codedError{
                                 m:    "transformed: " + e.m,
                                 code: "transformed: " + e.code,
                         }
                 }
@@ skipping 17 matching lines @@
         })
 
         client.Close()
         err = chanReadError(c, srvDone, "server done")
         c.Assert(err, IsNil)
 }
 
 func (*suite) TestServerWaitsForOutstandingCalls(c *C) {
         ready := make(chan struct{})
         start := make(chan string)
-        root := &TRoot{
+        root := &Root{
                 delayed: map[string]*DelayedMethods{
                         "1": {
                                 ready: ready,
                                 done:  start,
                         },
                 },
         }
         client, srvDone := newRPCClientServer(c, root, nil)
         done := make(chan struct{})
         go func() {
@@ skipping 19 matching lines @@
 func chanRead(c *C, ch <-chan struct{}, what string) {
         select {
         case <-ch:
                 return
         case <-time.After(3 * time.Second):
                 c.Fatalf("timeout on channel read %s", what)
         }
 }
 
 func (*suite) TestCompatibility(c *C) {
-        root := &TRoot{
+        root := &Root{
                 simple: make(map[string]*SimpleMethods),
         }
         a0 := &SimpleMethods{root: root, id: "a0"}
         root.simple["a0"] = a0
 
         client, srvDone := newRPCClientServer(c, root, nil)
         call := func(method string, arg, ret interface{}) (passedArg interface{}) {
                 root.calls = nil
                 err := client.Call("SimpleMethods", "a0", method, arg, ret)
                 c.Assert(err, IsNil)
@@ skipping 26 matching lines @@
         arg = call("Call1r0", stringVal{"x"}, &r)
         c.Assert(arg, Equals, stringVal{"x"})
         c.Assert(r, Equals, extra{})
 
         client.Close()
         err := chanReadError(c, srvDone, "server done")
         c.Assert(err, IsNil)
 }
 
 func (*suite) TestBadCall(c *C) {
-        root := &TRoot{
+        root := &Root{
                 simple: make(map[string]*SimpleMethods),
         }
         a0 := &SimpleMethods{root: root, id: "a0"}
         root.simple["a0"] = a0
         client, srvDone := newRPCClientServer(c, root, nil)
 
         err := client.Call("BadSomething", "a0", "No", nil, nil)
         c.Assert(err, ErrorMatches, `server error: unknown object type "BadSomething"`)
 
         err = client.Call("SimpleMethods", "xx", "No", nil, nil)
         c.Assert(err, ErrorMatches, `server error: no such request "No" on SimpleMethods`)
 
         err = client.Call("SimpleMethods", "xx", "Call0r0", nil, nil)
         c.Assert(err, ErrorMatches, "server error: unknown SimpleMethods id")
 
         client.Close()
         err = chanReadError(c, srvDone, "server done")
         c.Assert(err, IsNil)
 }
 
 func (*suite) TestErrorAfterClientClose(c *C) {
-        client, srvDone := newRPCClientServer(c, &TRoot{}, nil)
+        client, srvDone := newRPCClientServer(c, &Root{}, nil)
         err := client.Close()
         c.Assert(err, IsNil)
         err = client.Call("Foo", "", "Bar", nil, nil)
         c.Assert(err, Equals, rpc.ErrShutdown)
         err = chanReadError(c, srvDone, "server done")
         c.Assert(err, IsNil)
 }
 
+type KillerRoot struct {
+        killed bool
+        Root
+}
+
+func (r *KillerRoot) Kill() {
+        r.killed = true
+}
+
+func (*suite) TestRootIsKilled(c *C) {
+        root := &KillerRoot{}
+        client, srvDone := newRPCClientServer(c, root, nil)
+        err := client.Close()
+        c.Assert(err, IsNil)
+        err = chanReadError(c, srvDone, "server done")
+        c.Assert(err, IsNil)
+        c.Assert(root.killed, Equals, true)
+}
+
 func chanReadError(c *C, ch <-chan error, what string) error {
         select {
         case e := <-ch:
                 return e
         case <-time.After(3 * time.Second):
                 c.Fatalf("timeout on channel read %s", what)
         }
         panic("unreachable")
 }
 
 // newRPCClientServer starts an RPC server serving a connection from a
 // single client.  When the server has finished serving the connection,
 // it sends a value on done.
 func newRPCClientServer(c *C, root interface{}, tfErr func(error) error) (client *rpc.Client, done <-chan error) {
-        srv, err := rpc.NewServer(&TRoot{}, tfErr)
+        srv, err := rpc.NewServer(root, tfErr)
         c.Assert(err, IsNil)
 
         l, err := net.Listen("tcp", ":0")
         c.Assert(err, IsNil)
         defer l.Close()
 
         srvDone := make(chan error, 1)
         go func() {
                 conn, err := l.Accept()
                 if err != nil {
@@ skipping 92 matching lines @@
         }
 }
 
 func NewJSONClientCodec(c io.ReadWriteCloser) rpc.ClientCodec {
         return &generalClientCodec{
                 Closer: c,
                 enc:    json.NewEncoder(c),
                 dec:    json.NewDecoder(c),
         }
 }