Skip to content

Commit

Permalink
remove goroutines from linear reader
Browse files Browse the repository at this point in the history
  • Loading branch information
divyam234 committed Nov 2, 2023
1 parent 2533234 commit 6111362
Showing 1 changed file with 45 additions and 79 deletions.
124 changes: 45 additions & 79 deletions utils/reader/lr.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,17 @@ type linearReader struct {
ctx context.Context
parts []types.Part
pos int
reader io.ReadCloser
client *telegram.Client
next func() []byte
buffer []byte
bytesread int64
chunkSize int64
sync.Mutex
i int64
mu sync.Mutex
}

func (*linearReader) Close() error {
return nil
}

func NewLinearReader(ctx context.Context, client *telegram.Client, parts []types.Part) (io.ReadCloser, error) {
Expand All @@ -31,47 +37,37 @@ func NewLinearReader(ctx context.Context, client *telegram.Client, parts []types
chunkSize: int64(1024 * 1024),
}

res, err := r.nextPart()

if err != nil {
return nil, err
}

r.reader = res
r.next = r.partStream()

return r, nil
}

func (r *linearReader) Read(p []byte) (n int, err error) {
r.Lock()
defer r.Unlock()
n, err = r.reader.Read(p)
if err != nil {
return 0, err
r.mu.Lock()
defer r.mu.Unlock()

if r.i >= int64(len(r.buffer)) {
r.buffer = r.next()
if len(r.buffer) == 0 && r.pos == len(r.parts)-1 {
return 0, io.EOF
}
r.i = 0
}

n = copy(p, r.buffer[r.i:])

r.i += int64(n)

r.bytesread += int64(n)

if r.bytesread == r.parts[r.pos].Length && r.pos < len(r.parts)-1 {
r.pos++
r.reader, err = r.nextPart()

if err != nil {
return 0, err
}
r.next = r.partStream()
r.bytesread = 0
}
return n, nil
}

func (r *linearReader) Close() (err error) {
if r.reader != nil {
err = r.reader.Close()
r.reader = nil
}
return
}

func (r *linearReader) chunk(offset int64, limit int64) ([]byte, error) {

req := &tg.UploadGetFileRequest{
Expand All @@ -94,31 +90,9 @@ func (r *linearReader) chunk(offset int64, limit int64) ([]byte, error) {
}
}

func (r *linearReader) nextPart() (io.ReadCloser, error) {
stream := r.tgRangeStream()
ir, iw := io.Pipe()

go func() {
defer iw.Close()

for {

data, more := <-stream
if !more {
return
}

_, err := iw.Write(data)
if err != nil {
return
}
}
}()

return ir, nil
}

func (r *linearReader) tgRangeStream() chan []byte {
func (r *linearReader) partStream() func() []byte {
r.mu.Lock()
defer r.mu.Unlock()

start := r.parts[r.pos].Start
end := r.parts[r.pos].End
Expand All @@ -132,41 +106,33 @@ func (r *linearReader) tgRangeStream() chan []byte {

currentPart := 1

stream := make(chan []byte)
readData := func() []byte {

go func() {

defer close(stream)

for {

res, _ := r.chunk(offset, r.chunkSize)

if len(res) == 0 {
return
} else if partCount == 1 {
res = res[firstPartCut:lastPartCut]
if currentPart > partCount {
return make([]byte, 0)
}

} else if currentPart == 1 {
res = res[firstPartCut:]
res, _ := r.chunk(offset, r.chunkSize)

} else if currentPart == partCount {
res = res[:lastPartCut]
if len(res) == 0 {
return res
} else if partCount == 1 {
res = res[firstPartCut:lastPartCut]

}
} else if currentPart == 1 {
res = res[firstPartCut:]

stream <- res
} else if currentPart == partCount {
res = res[:lastPartCut]

currentPart++
}

offset += r.chunkSize
currentPart++

if currentPart > partCount {
return
}
offset += r.chunkSize

}
}()
return res

return stream
}
return readData
}

0 comments on commit 6111362

Please sign in to comment.