/
ctxcopy.go
79 lines (67 loc) · 1.83 KB
/
ctxcopy.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
package ctxcopy
import (
"fmt"
"io"
"golang.org/x/net/context"
)
const (
// zeroBufferLenError is error message of buffer length.
zeroBufferLenError string = "Buffer length is 0."
// canceledStr is the message of copy operation canceled.
canceledStr string = "Copy operation is canceled."
)
// copyWithCancelation reads bytes to buffer from source and write to destination from buffer with cancelation signal.
//
// Params:
// dst: Destination.
// src: Source.
// buf: Buffer(length should >= 0).
// cancel: cancelation flag. It'll stop copy while the cancelation signal is set to true.
func copyWithCancelation(dst io.Writer, src io.Reader, buf []byte, cancel *bool) (err error) {
if len(buf) == 0 {
return fmt.Errorf(zeroBufferLenError)
}
for {
if cancel != nil && (*cancel) {
return fmt.Errorf(canceledStr)
}
n, err := src.Read(buf)
if err != nil && err != io.EOF {
return err
}
if n == 0 {
break
}
if _, err := dst.Write(buf[:n]); err != nil {
return err
}
}
return nil
}
// Copy reads bytes to buffer from source and write to destination from buffer with cancelation signal.
//
// Params:
// ctx:
// Google's Context type which carries deadlines, cacelation signals,
// and other request-scoped values across API boundaries and between processes.
// See https://godoc.org/golang.org/x/net/context for more.
// dst: Destination.
// src: Source.
// buf: Buffer(length should >= 0).
func Copy(ctx context.Context, dst io.Writer, src io.Reader, buf []byte) (err error) {
cancel := false
c := make(chan error)
go func() {
c <- copyWithCancelation(dst, src, buf, &cancel)
}()
select {
case err = <-c:
return err
case <-ctx.Done():
// Set cancelation flag to true
cancel = true
// Wait for copyWithCancelation() return.
<-c
return ctx.Err()
}
}