流式传输到fasthttp.RequestCtx

时间:2018-07-30 15:08:29

标签: go fasthttp

我想让客户端执行流API并使用fasthttp作为后端。 我不想向客户端公开fasthttp.RequestCtx,也不想从客户端功能获取频道。

这是我的代码(客户端函数是worker),当我运行它然后执行curl -iv http://localhost:8080时,它卡住了。在服务器端,我看到5张WRITE: ...打印件。调试似乎服务器挂在s.writer.Flush()

任何想法如何实现这一目标?

package main

import (
    "bufio"
    "fmt"
    "log"

    "github.com/valyala/fasthttp"
)

var (
    poem = []string{
        "The Road goes ever on and on",
        "Down from the door where it began.",
        "Now far ahead the Road has gone,",
        "And I must follow, if I can,",
        "Pursuing it with eager feet,",
        "Until it joins some larger way",
        "Where many paths and errands meet.",
        "And whither then? I cannot say.",
    }
)

// Streamer is streaming API
type Streamer interface {
    // Write writes bytes to streamer
    Write(p []byte) (n int, err error)
    // Flush flushes data to the client
    Flush() error
    // SetStatusCode sets the status code. *Must* be called before Write and Flush
    SetStatusCode(statusCode int) error
    // SetHeader sets a response header. *Must* be called before Write and Flush
    SetHeader(key string, value interface{}) error
}

// HTTPStreamer implement Streamer with fasthttp
type HTTPStreamer struct {
    ctx    *fasthttp.RequestCtx
    done   chan bool
    writer *bufio.Writer
}

// NewHTTPStreamer returns a new HTTPStreamer
func NewHTTPStreamer(ctx *fasthttp.RequestCtx) *HTTPStreamer {
    return &HTTPStreamer{
        ctx: ctx,
    }
}

// Close closes the stream
func (s *HTTPStreamer) Close() {
    if s.done != nil {
        close(s.done)
    }
}

// Write writes bytes to streamer
func (s *HTTPStreamer) Write(p []byte) (n int, err error) {
    if s.writer == nil {
        s.done = make(chan bool)
        ready := make(chan bool)
        go func() {
            s.ctx.SetBodyStreamWriter(func(w *bufio.Writer) {
                s.writer = w
                ready <- true // Signal that writer is set
                <-s.done      // Wait for stream to be closed
            })
        }()

        <-ready // Wait until writer is set
    }

    fmt.Printf("WRITE: %s", string(p))
    return s.writer.Write(p)
}

// Flush flushes data to the client
func (s *HTTPStreamer) Flush() error {
    if s.writer != nil {
        return s.writer.Flush()
    }

    return nil
}

// SetStatusCode sets the status code. *Must* be called before Write and Flush
func (s *HTTPStreamer) SetStatusCode(statusCode int) error {
    if s.writer != nil {
        return fmt.Errorf("Streaming started - can't set status")
    }

    s.ctx.SetStatusCode(statusCode)
    return nil
}

// SetHeader sets a response header. *Must* be called before Write and Flush
// value can be string or []byte
func (s *HTTPStreamer) SetHeader(key string, value interface{}) error {
    if s.writer != nil {
        return fmt.Errorf("Streaming started - can't set header")
    }

    switch v := value.(type) {
    case string:
        s.ctx.Response.Header.Set(key, v)
    case []byte:
        s.ctx.Response.Header.SetBytesV(key, v)
    default:
        return fmt.Errorf("Unsupported header value type - %T", value)
    }

    return nil
}

func worker(s Streamer) {
    s.SetStatusCode(201)
    s.SetHeader("X-T", "VALUE")
    for _, line := range poem {
        s.Write([]byte(line + "\n"))
        s.Flush()
    }
}

func handler(ctx *fasthttp.RequestCtx) {
    s := NewHTTPStreamer(ctx)
    worker(s)
    s.Close()
}

func main() {
    srv := &fasthttp.Server{
        Handler: handler,
    }

    if err := srv.ListenAndServe(":8080"); err != nil {
        log.Fatal(err)
    }
}

0 个答案:

没有答案