Golang rpc.Message類代碼示例

func (p *pipeTransport) SendMessage(ctx context.Context, msg rpccapnp.Message) error {
	if !p.startSend() {
		return errClosed
	}
	defer p.finishSend()

	buf, err := msg.Segment().Message().Marshal()
	if err != nil {
		return err
	}
	mm, err := capnp.Unmarshal(buf)
	if err != nil {
		return err
	}
	msg, err = rpccapnp.ReadRootMessage(mm)
	if err != nil {
		return err
	}

	select {
	case p.w <- msg:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	case <-p.finish:
		return errClosed
	case <-p.otherFin:
		return errBrokenPipe
	}
}

// copyRPCMessage clones an RPC packet.
func copyRPCMessage(m rpccapnp.Message) rpccapnp.Message {
	mm := copyMessage(m.Segment().Message())
	rpcMsg, err := rpccapnp.ReadRootMessage(mm)
	if err != nil {
		panic(err)
	}
	return rpcMsg
}

// teardown moves the connection from the dying to the dead state.
func (c *Conn) teardown(abort rpccapnp.Message) {
	c.workers.Wait()

	c.mu.Lock()
	for _, q := range c.questions {
		if q != nil {
			q.cancel(ErrConnClosed)
		}
	}
	c.questions = nil
	exps := c.exports
	c.exports = nil
	c.embargoes = nil
	for _, a := range c.answers {
		a.cancel()
	}
	c.answers = nil
	c.imports = nil
	c.mainFunc = nil
	c.mu.Unlock()

	if c.mainCloser != nil {
		if err := c.mainCloser.Close(); err != nil {
			c.errorf("closing main interface: %v", err)
		}
		c.mainCloser = nil
	}
	// Closing an export may try to lock the Conn, so run it outside
	// critical section.
	for id, e := range exps {
		if e == nil {
			continue
		}
		if err := e.client.Close(); err != nil {
			c.errorf("export %v close: %v", id, err)
		}
	}
	exps = nil

	var werr error
	if abort.IsValid() {
		werr = c.transport.SendMessage(context.Background(), abort)
	}
	cerr := c.transport.Close()

	c.stateMu.Lock()
	if c.closeErr == ErrConnClosed {
		if cerr != nil {
			c.closeErr = cerr
		} else if werr != nil {
			c.closeErr = werr
		}
	}
	c.state = connDead
	c.stateCond.Broadcast()
	c.stateMu.Unlock()
}

func (c *Conn) handleDisembargoMessage(msg rpccapnp.Message) error {
	d, err := msg.Disembargo()
	if err != nil {
		return err
	}
	dtarget, err := d.Target()
	if err != nil {
		return err
	}
	switch d.Context().Which() {
	case rpccapnp.Disembargo_context_Which_senderLoopback:
		id := embargoID(d.Context().SenderLoopback())
		if dtarget.Which() != rpccapnp.MessageTarget_Which_promisedAnswer {
			return errDisembargoNonImport
		}
		dpa, err := dtarget.PromisedAnswer()
		if err != nil {
			return err
		}
		aid := answerID(dpa.QuestionId())
		a := c.answers[aid]
		if a == nil {
			return errDisembargoMissingAnswer
		}
		dtrans, err := dpa.Transform()
		if err != nil {
			return err
		}
		transform := promisedAnswerOpsToTransform(dtrans)
		queued, err := a.queueDisembargo(transform, id, dtarget)
		if err != nil {
			return err
		}
		if !queued {
			// There's nothing to embargo; everything's been delivered.
			resp := newDisembargoMessage(nil, rpccapnp.Disembargo_context_Which_receiverLoopback, id)
			rd, _ := resp.Disembargo()
			if err := rd.SetTarget(dtarget); err != nil {
				return err
			}
			c.sendMessage(resp)
		}
	case rpccapnp.Disembargo_context_Which_receiverLoopback:
		id := embargoID(d.Context().ReceiverLoopback())
		c.disembargo(id)
	default:
		um := newUnimplementedMessage(nil, msg)
		c.sendMessage(um)
	}
	return nil
}

// handleCallMessage handles a received call message.  It mutates the
// capability table of its parameter.  The caller holds onto c.mu.
func (c *Conn) handleCallMessage(m rpccapnp.Message) error {
	mcall, err := m.Call()
	if err != nil {
		return err
	}
	mt, err := mcall.Target()
	if err != nil {
		return err
	}
	if mt.Which() != rpccapnp.MessageTarget_Which_importedCap && mt.Which() != rpccapnp.MessageTarget_Which_promisedAnswer {
		um := newUnimplementedMessage(nil, m)
		return c.sendMessage(um)
	}
	mparams, err := mcall.Params()
	if err != nil {
		return err
	}
	if err := c.populateMessageCapTable(mparams); err == errUnimplemented {
		um := newUnimplementedMessage(nil, m)
		return c.sendMessage(um)
	} else if err != nil {
		c.abort(err)
		return err
	}
	ctx, cancel := c.newContext()
	id := answerID(mcall.QuestionId())
	a := c.insertAnswer(id, cancel)
	if a == nil {
		// Question ID reused, error out.
		c.abort(errQuestionReused)
		return errQuestionReused
	}
	meth := capnp.Method{
		InterfaceID: mcall.InterfaceId(),
		MethodID:    mcall.MethodId(),
	}
	paramContent, err := mparams.ContentPtr()
	if err != nil {
		return err
	}
	cl := &capnp.Call{
		Ctx:    ctx,
		Method: meth,
		Params: paramContent.Struct(),
	}
	if err := c.routeCallMessage(a, mt, cl); err != nil {
		return a.reject(err)
	}
	return nil
}

func copyAbort(m rpccapnp.Message) (Abort, error) {
	ma, err := m.Abort()
	if err != nil {
		return Abort{}, err
	}
	msg, _, _ := capnp.NewMessage(capnp.SingleSegment(nil))
	if err := msg.SetRootPtr(ma.ToPtr()); err != nil {
		return Abort{}, err
	}
	p, err := msg.RootPtr()
	if err != nil {
		return Abort{}, err
	}
	return Abort{rpccapnp.Exception{Struct: p.Struct()}}, nil
}

func (s *streamTransport) SendMessage(ctx context.Context, msg rpccapnp.Message) error {
	s.wbuf.Reset()
	if err := s.enc.Encode(msg.Segment().Message()); err != nil {
		return err
	}
	if s.deadline != nil {
		// TODO(light): log errors
		if d, ok := ctx.Deadline(); ok {
			s.deadline.SetWriteDeadline(d)
		} else {
			s.deadline.SetWriteDeadline(time.Time{})
		}
	}
	_, err := s.rwc.Write(s.wbuf.Bytes())
	return err
}

// handleReturnMessage is to handle a received return message.
// The caller is holding onto c.mu.
func (c *Conn) handleReturnMessage(m rpccapnp.Message) error {
	ret, err := m.Return()
	if err != nil {
		return err
	}
	id := questionID(ret.AnswerId())
	q := c.popQuestion(id)
	if q == nil {
		return fmt.Errorf("received return for unknown question id=%d", id)
	}
	if ret.ReleaseParamCaps() {
		for _, id := range q.paramCaps {
			c.releaseExport(id, 1)
		}
	}
	q.mu.RLock()
	qstate := q.state
	q.mu.RUnlock()
	if qstate == questionCanceled {
		// We already sent the finish message.
		return nil
	}
	releaseResultCaps := true
	switch ret.Which() {
	case rpccapnp.Return_Which_results:
		releaseResultCaps = false
		results, err := ret.Results()
		if err != nil {
			return err
		}
		if err := c.populateMessageCapTable(results); err == errUnimplemented {
			um := newUnimplementedMessage(nil, m)
			c.sendMessage(um)
			return errUnimplemented
		} else if err != nil {
			c.abort(err)
			return err
		}
		content, err := results.ContentPtr()
		if err != nil {
			return err
		}
		q.fulfill(content)
	case rpccapnp.Return_Which_exception:
		exc, err := ret.Exception()
		if err != nil {
			return err
		}
		e := error(Exception{exc})
		if q.method != nil {
			e = &capnp.MethodError{
				Method: q.method,
				Err:    e,
			}
		} else {
			e = bootstrapError{e}
		}
		q.reject(e)
	case rpccapnp.Return_Which_canceled:
		err := &questionError{
			id:     id,
			method: q.method,
			err:    fmt.Errorf("receiver reported canceled"),
		}
		c.errorf("%v", err)
		q.reject(err)
		return nil
	default:
		um := newUnimplementedMessage(nil, m)
		c.sendMessage(um)
		return errUnimplemented
	}
	fin := newFinishMessage(nil, id, releaseResultCaps)
	c.sendMessage(fin)
	return nil
}

// handleMessage is run from the receive goroutine to process a single
// message.  m cannot be held onto past the return of handleMessage, and
// c.mu is not held at the start of handleMessage.
func (c *Conn) handleMessage(m rpccapnp.Message) {
	switch m.Which() {
	case rpccapnp.Message_Which_unimplemented:
		// no-op for now to avoid feedback loop
	case rpccapnp.Message_Which_abort:
		a, err := copyAbort(m)
		if err != nil {
			c.errorf("decode abort: %v", err)
			// Keep going, since we're trying to abort anyway.
		}
		c.infof("abort: %v", a)
		c.shutdown(a)
	case rpccapnp.Message_Which_return:
		m = copyRPCMessage(m)
		c.mu.Lock()
		err := c.handleReturnMessage(m)
		c.mu.Unlock()

		if err != nil {
			c.errorf("handle return: %v", err)
		}
	case rpccapnp.Message_Which_finish:
		mfin, err := m.Finish()
		if err != nil {
			c.errorf("decode finish: %v", err)
			return
		}
		id := answerID(mfin.QuestionId())

		c.mu.Lock()
		a := c.popAnswer(id)
		if a == nil {
			c.mu.Unlock()
			c.errorf("finish called for unknown answer %d", id)
			return
		}
		a.cancel()
		if mfin.ReleaseResultCaps() {
			for _, id := range a.resultCaps {
				c.releaseExport(id, 1)
			}
		}
		c.mu.Unlock()
	case rpccapnp.Message_Which_bootstrap:
		boot, err := m.Bootstrap()
		if err != nil {
			c.errorf("decode bootstrap: %v", err)
			return
		}
		id := answerID(boot.QuestionId())

		c.mu.Lock()
		err = c.handleBootstrapMessage(id)
		c.mu.Unlock()

		if err != nil {
			c.errorf("handle bootstrap: %v", err)
		}
	case rpccapnp.Message_Which_call:
		m = copyRPCMessage(m)
		c.mu.Lock()
		err := c.handleCallMessage(m)
		c.mu.Unlock()

		if err != nil {
			c.errorf("handle call: %v", err)
		}
	case rpccapnp.Message_Which_release:
		rel, err := m.Release()
		if err != nil {
			c.errorf("decode release: %v", err)
			return
		}
		id := exportID(rel.Id())
		refs := int(rel.ReferenceCount())

		c.mu.Lock()
		c.releaseExport(id, refs)
		c.mu.Unlock()
	case rpccapnp.Message_Which_disembargo:
		m = copyRPCMessage(m)
		c.mu.Lock()
		err := c.handleDisembargoMessage(m)
		c.mu.Unlock()

		if err != nil {
			// Any failure in a disembargo is a protocol violation.
			c.abort(err)
		}
	default:
		c.infof("received unimplemented message, which = %v", m.Which())
		um := newUnimplementedMessage(nil, m)
		c.sendMessage(um)
	}
}