一文带你搞懂go中的请求超时控制

func hardWork(job interface{}) error {
    time.Sleep(time.Minute)
    return nil
}

func requestWorkV1(ctx context.Context, job interface{}) error {
    ctx, cancel := context.WithTimeout(ctx, time.Second*2)
    defer cancel()

    // 仅需要改这里即可
    // done := make(chan error, 1)
    done := make(chan error)

    // done 退出以后，没有接受者，会导致协程阻塞
    go func() {
        done <- hardWork(job)
    }()

    select {
        case err := <-done:
        return err
        case <-ctx.Done():   // 这一部分提前退出
        return ctx.Err()
    }
}

// 可以做到超时控制，但是会出现协程泄露的情况
func TestV1(t *testing.T) {
    const total = 1000
    var wg sync.WaitGroup
    wg.Add(total)
    now := time.Now()

    for i := 0; i < total; i++ {
        go func() {
            defer wg.Done()
            requestWorkV1(context.Background(), "any")
        }()
    }

    wg.Wait()
    fmt.Println("elapsed:", time.Since(now))  // 2秒后打印这条语句，说明协程只执行了两秒

    time.Sleep(time.Minute * 2)
    fmt.Println("number of goroutines:", runtime.NumGoroutine())  // number of goroutines: 1002
}

syntax = "proto3";

package helloworld;

option go_package = ".";

service Greeter {
  rpc SayHello (HelloRequest) returns (HelloReply);
}

message HelloRequest {
  string name = 1;
}

message HelloReply {
  string message = 1;
}

const (
    address     = "localhost:50051"
    defaultName = "world"
)

func TestClient(t *testing.T) {
    conn, err := grpc.Dial(address, grpc.WithInsecure())
    if err != nil {
        log.Fatalf("did not connect: %v", err)
    }
    defer conn.Close()
    c := pb.NewGreeterClient(conn)

    name := defaultName

    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()
    r, err := c.SayHello(ctx, &pb.HelloRequest{Name: name})
    if err != nil {
        log.Fatalf("could not greet: %v", err)
    }
    log.Printf("Greeting: %s", r.Message)
}

func (s *server) SayHello(ctx context.Context, request *pb.HelloRequest) (*pb.HelloReply, error) {
    logrus.Info("request in")

    time.Sleep(5 * time.Second)

    //select {
    //case <-ctx.Done():
    // fmt.Println("time out Done")
    //}

    logrus.Info("requst out")

    if ctx.Err() == context.DeadlineExceeded {
        log.Printf("RPC has reached deadline exceeded state: %s", ctx.Err())
        return nil, ctx.Err()
    }

    return &pb.HelloReply{Message: "Hello, " + request.Name}, nil
}

func TestServer(t *testing.T) {
    lis, err := net.Listen("tcp", ":50051")
    if err != nil {
        log.Fatalf("Failed to listen: %v", err)
    }

    s := grpc.NewServer()
    pb.RegisterGreeterServer(s, &server{})
    if err := s.Serve(lis); err != nil {
        log.Fatalf("Failed to serve: %v", err)
    }
}

func newClientStreamWithParams(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, mc serviceconfig.MethodConfig, onCommit, doneFunc func(), opts ...CallOption) (_ iresolver.ClientStream, err error) {
    // .....


	if desc != unaryStreamDesc {
		// Listen on cc and stream contexts to cleanup when the user closes the
		// ClientConn or cancels the stream context.  In all other cases, an error
		// should already be injected into the recv buffer by the transport, which
		// the client will eventually receive, and then we will cancel the stream's
		// context in clientStream.finish.
		go func() {
			select {
			case <-cc.ctx.Done():
				cs.finish(ErrClientConnClosing)
			case <-ctx.Done():
				cs.finish(toRPCErr(ctx.Err()))
			}
		}()
	}
}