cancelreader_select.go 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. //go:build solaris || darwin || freebsd || netbsd || openbsd || dragonfly
  2. // +build solaris darwin freebsd netbsd openbsd dragonfly
  3. package cancelreader
  4. import (
  5. "errors"
  6. "fmt"
  7. "io"
  8. "os"
  9. "strings"
  10. "golang.org/x/sys/unix"
  11. )
  12. // newSelectCancelReader returns a reader and a cancel function. If the input
  13. // reader is a File, the cancel function can be used to interrupt a
  14. // blocking call read call. In this case, the cancel function returns true if
  15. // the call was canceled successfully. If the input reader is not a File or
  16. // the file descriptor is 1024 or larger, the cancel function does nothing and
  17. // always returns false. The generic unix implementation is based on the posix
  18. // select syscall.
  19. func newSelectCancelReader(reader io.Reader) (CancelReader, error) {
  20. file, ok := reader.(File)
  21. if !ok || file.Fd() >= unix.FD_SETSIZE {
  22. return newFallbackCancelReader(reader)
  23. }
  24. r := &selectCancelReader{file: file}
  25. var err error
  26. r.cancelSignalReader, r.cancelSignalWriter, err = os.Pipe()
  27. if err != nil {
  28. return nil, err
  29. }
  30. return r, nil
  31. }
  32. type selectCancelReader struct {
  33. file File
  34. cancelSignalReader File
  35. cancelSignalWriter File
  36. cancelMixin
  37. }
  38. func (r *selectCancelReader) Read(data []byte) (int, error) {
  39. if r.isCanceled() {
  40. return 0, ErrCanceled
  41. }
  42. for {
  43. err := waitForRead(r.file, r.cancelSignalReader)
  44. if err != nil {
  45. if errors.Is(err, unix.EINTR) {
  46. continue // try again if the syscall was interrupted
  47. }
  48. if errors.Is(err, ErrCanceled) {
  49. // remove signal from pipe
  50. var b [1]byte
  51. _, readErr := r.cancelSignalReader.Read(b[:])
  52. if readErr != nil {
  53. return 0, fmt.Errorf("reading cancel signal: %w", readErr)
  54. }
  55. }
  56. return 0, err
  57. }
  58. return r.file.Read(data)
  59. }
  60. }
  61. func (r *selectCancelReader) Cancel() bool {
  62. r.setCanceled()
  63. // send cancel signal
  64. _, err := r.cancelSignalWriter.Write([]byte{'c'})
  65. return err == nil
  66. }
  67. func (r *selectCancelReader) Close() error {
  68. var errMsgs []string
  69. // close pipe
  70. err := r.cancelSignalWriter.Close()
  71. if err != nil {
  72. errMsgs = append(errMsgs, fmt.Sprintf("closing cancel signal writer: %v", err))
  73. }
  74. err = r.cancelSignalReader.Close()
  75. if err != nil {
  76. errMsgs = append(errMsgs, fmt.Sprintf("closing cancel signal reader: %v", err))
  77. }
  78. if len(errMsgs) > 0 {
  79. return fmt.Errorf(strings.Join(errMsgs, ", "))
  80. }
  81. return nil
  82. }
  83. func waitForRead(reader, abort File) error {
  84. readerFd := int(reader.Fd())
  85. abortFd := int(abort.Fd())
  86. maxFd := readerFd
  87. if abortFd > maxFd {
  88. maxFd = abortFd
  89. }
  90. // this is a limitation of the select syscall
  91. if maxFd >= unix.FD_SETSIZE {
  92. return fmt.Errorf("cannot select on file descriptor %d which is larger than 1024", maxFd)
  93. }
  94. fdSet := &unix.FdSet{}
  95. fdSet.Set(int(reader.Fd()))
  96. fdSet.Set(int(abort.Fd()))
  97. _, err := unix.Select(maxFd+1, fdSet, nil, nil, nil)
  98. if err != nil {
  99. return fmt.Errorf("select: %w", err)
  100. }
  101. if fdSet.IsSet(abortFd) {
  102. return ErrCanceled
  103. }
  104. if fdSet.IsSet(readerFd) {
  105. return nil
  106. }
  107. return fmt.Errorf("select returned without setting a file descriptor")
  108. }