Exemplo n.º 1
0
func (lp *leaseProxy) LeaseKeepAlive(stream pb.Lease_LeaseKeepAliveServer) error {
	conn := lp.client.ActiveConnection()
	ctx, cancel := context.WithCancel(stream.Context())
	lc, err := pb.NewLeaseClient(conn).LeaseKeepAlive(ctx)
	if err != nil {
		cancel()
		return err
	}

	go func() {
		// Cancel the context attached to lc to unblock lc.Recv when
		// this routine returns on error.
		defer cancel()

		for {
			// stream.Recv will be unblock when the loop in the parent routine
			// returns on error.
			rr, err := stream.Recv()
			if err != nil {
				return
			}
			err = lc.Send(rr)
			if err != nil {
				return
			}
		}
	}()

	for {
		rr, err := lc.Recv()
		if err != nil {
			return err
		}
		err = stream.Send(rr)
		if err != nil {
			return err
		}
	}
}
Exemplo n.º 2
0
func (ls *LeaseServer) LeaseKeepAlive(stream pb.Lease_LeaseKeepAliveServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}

		// Create header before we sent out the renew request.
		// This can make sure that the revision is strictly smaller or equal to
		// when the keepalive happened at the local server (when the local server is the leader)
		// or remote leader.
		// Without this, a lease might be revoked at rev 3 but client can see the keepalive succeeded
		// at rev 4.
		resp := &pb.LeaseKeepAliveResponse{ID: req.ID, Header: &pb.ResponseHeader{}}
		ls.hdr.fill(resp.Header)

		ttl, err := ls.le.LeaseRenew(stream.Context(), lease.LeaseID(req.ID))
		if err == lease.ErrLeaseNotFound {
			err = nil
			ttl = 0
		}

		if err != nil {
			return togRPCError(err)
		}

		resp.TTL = ttl
		err = stream.Send(resp)
		if err != nil {
			return err
		}
	}
}