github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/iokit/traverse.go (about)

     1  // Copyright 2020 Insolar Network Ltd.
     2  // All rights reserved.
     3  // This material is licensed under the Insolar License version 1.0,
     4  // available at https://github.com/insolar/assured-ledger/blob/master/LICENSE.md.
     5  
     6  package iokit
     7  
     8  import (
     9  	"bytes"
    10  	"io"
    11  	"math"
    12  	"strings"
    13  )
    14  
    15  func LimitOfReader(r io.Reader) int64 {
    16  	_, n := traverseLimitReaders(r, math.MaxInt64)
    17  	return n
    18  }
    19  
    20  func LimitOfWriter(w io.Writer) int64 {
    21  	_, n := traverseLimitWriters(w, math.MaxInt64)
    22  	return n
    23  }
    24  
    25  func LimitOfWriterTo(w io.WriterTo) int64 {
    26  	switch ww := w.(type) {
    27  	case interface{ FixedByteSize() int }:
    28  		return int64(ww.FixedByteSize())
    29  	case interface{ ProtoSize() int }:
    30  		return int64(ww.ProtoSize())
    31  	case interface{ Size() int }:
    32  		return int64(ww.Size())
    33  	case io.Reader:
    34  		if n := sizeOfReader(ww); n >= 0 {
    35  			return n
    36  		}
    37  	}
    38  	return math.MaxInt64
    39  }
    40  
    41  func sizeOfReader(r io.Reader) int64 {
    42  	switch rr := r.(type) {
    43  	case *bytes.Buffer:
    44  		return int64(rr.Len())
    45  	case *strings.Reader:
    46  		return int64(rr.Len())
    47  	case interface{ PhysicalRemaining() int64 }:
    48  		return rr.PhysicalRemaining()
    49  	//case interface{ Len() int }:
    50  	//	return int64(rr.Len())
    51  	default:
    52  		return -1
    53  	}
    54  }
    55  
    56  func traverseLimitReaders(r io.Reader, min int64) (io.Reader, int64) {
    57  	for r != nil {
    58  		if min <= 0 {
    59  			return r, 0
    60  		}
    61  
    62  		var n int64
    63  		switch v := r.(type) {
    64  		case *RateLimitedReader:
    65  			_, min = traverseLimitReaders(v.R, min)
    66  			return r, min
    67  		case *TeeReader:
    68  			if v.CopyTo == nil {
    69  				r = v.R
    70  				continue
    71  			}
    72  			_, min = traverseLimitReaders(v.R, min)
    73  			return v, min
    74  		case *io.LimitedReader:
    75  			n = v.N
    76  			r = v.R
    77  		case *LimitedTriggerReader:
    78  			r = &v.hLimitedReader
    79  			continue
    80  		case *LimitedReader:
    81  			n = v.n
    82  			r = &v.r
    83  		default:
    84  			switch n = sizeOfReader(r); {
    85  			case n >= 0 && n < min:
    86  				return r, n
    87  			default:
    88  				return r, min
    89  			}
    90  		}
    91  
    92  		if n < min {
    93  			min = n
    94  		}
    95  	}
    96  
    97  	return nil, 0
    98  }
    99  
   100  func updatedLimitReaders(r io.Reader, delta int64) {
   101  	if delta <= 0 {
   102  		return
   103  	}
   104  
   105  	for r != nil {
   106  		switch v := r.(type) {
   107  		case *TeeReader:
   108  			if v.CopyTo == nil {
   109  				r = v.R
   110  				continue
   111  			}
   112  			return
   113  		case *io.LimitedReader:
   114  			v.N -= delta
   115  			r = v.R
   116  		case *LimitedTriggerReader:
   117  			v.n -= delta
   118  			r = &v.hLimitedReader.r
   119  			if v.n <= 0 {
   120  				// this ensures proper sequence of triggering when there are
   121  				// multiple instances of LimitedTriggerReader on the read chain
   122  				updatedLimitReaders(r, delta)
   123  				v.trigger()
   124  				return
   125  			}
   126  		case *LimitedReader:
   127  			v.n -= delta
   128  			r = &v.r
   129  		default:
   130  			return
   131  		}
   132  	}
   133  }
   134  
   135  func traverseLimitWriters(w io.Writer, min int64) (io.Writer, int64) {
   136  	for w != nil {
   137  		if min <= 0 {
   138  			return w, 0
   139  		}
   140  
   141  		var n int64
   142  		switch v := w.(type) {
   143  		case *RateLimitedWriter:
   144  			_, min = traverseLimitWriters(v.W, min)
   145  			return w, min
   146  		case *TeeWriter:
   147  			if v.CopyTo == nil {
   148  				w = v.W
   149  				continue
   150  			}
   151  			_, min = traverseLimitWriters(v.W, min)
   152  			return v, min
   153  		case *LimitedWriter:
   154  			n = v.n
   155  			w = &v.w
   156  		default:
   157  			return w, min
   158  		}
   159  
   160  		if n < min {
   161  			min = n
   162  		}
   163  	}
   164  
   165  	return nil, 0
   166  }
   167  
   168  func updatedLimitWriters(w io.Writer, delta int64) {
   169  	if delta <= 0 {
   170  		return
   171  	}
   172  
   173  	for w != nil {
   174  		switch v := w.(type) {
   175  		case *TeeWriter:
   176  			if v.CopyTo == nil {
   177  				w = v.W
   178  				continue
   179  			}
   180  			return
   181  		case *LimitedWriter:
   182  			v.n -= delta
   183  			w = &v.w
   184  		default:
   185  			return
   186  		}
   187  	}
   188  }