response.go 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. package stream
  2. import (
  3. "errors"
  4. "sync"
  5. "github.com/gammazero/deque"
  6. )
  7. type StreamResponse[T any] struct {
  8. q deque.Deque[T]
  9. l *sync.Mutex
  10. sig chan bool
  11. closed bool
  12. max int
  13. listening bool
  14. onClose func()
  15. }
  16. func NewStreamResponse[T any](max int) *StreamResponse[T] {
  17. return &StreamResponse[T]{
  18. l: &sync.Mutex{},
  19. sig: make(chan bool),
  20. max: max,
  21. }
  22. }
  23. func (r *StreamResponse[T]) OnClose(f func()) {
  24. r.onClose = f
  25. }
  26. func (r *StreamResponse[T]) Next() bool {
  27. r.l.Lock()
  28. if r.closed {
  29. r.l.Unlock()
  30. return false
  31. }
  32. if r.q.Len() > 0 {
  33. r.l.Unlock()
  34. return true
  35. }
  36. r.listening = true
  37. defer func() {
  38. r.listening = false
  39. }()
  40. r.l.Unlock()
  41. return <-r.sig
  42. }
  43. func (r *StreamResponse[T]) Read() (T, error) {
  44. r.l.Lock()
  45. defer r.l.Unlock()
  46. if r.q.Len() > 0 {
  47. data := r.q.PopFront()
  48. return data, nil
  49. } else {
  50. var data T
  51. return data, errors.New("no data available, please call Next() to wait for data")
  52. }
  53. }
  54. func (r *StreamResponse[T]) Write(data T) error {
  55. r.l.Lock()
  56. if r.closed {
  57. r.l.Unlock()
  58. return nil
  59. }
  60. if r.q.Len() >= r.max {
  61. r.l.Unlock()
  62. return errors.New("queue is full")
  63. }
  64. r.q.PushBack(data)
  65. if r.q.Len() == 1 {
  66. if r.listening {
  67. r.sig <- true
  68. }
  69. }
  70. r.l.Unlock()
  71. return nil
  72. }
  73. func (r *StreamResponse[T]) Close() {
  74. r.l.Lock()
  75. if r.closed {
  76. r.l.Unlock()
  77. return
  78. }
  79. r.closed = true
  80. r.l.Unlock()
  81. select {
  82. case r.sig <- false:
  83. default:
  84. }
  85. close(r.sig)
  86. if r.onClose != nil {
  87. r.onClose()
  88. }
  89. }
  90. func (r *StreamResponse[T]) IsClosed() bool {
  91. r.l.Lock()
  92. defer r.l.Unlock()
  93. return r.closed
  94. }
  95. func (r *StreamResponse[T]) Size() int {
  96. r.l.Lock()
  97. defer r.l.Unlock()
  98. return r.q.Len()
  99. }