github.com/dshekhar95/sub_dgraph@v0.0.0-20230424164411-6be28e40bbf1/dgraph/cmd/bulk/progress.go (about)

     1  /*
     2   * Copyright 2017-2022 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package bulk
    18  
    19  import (
    20  	"fmt"
    21  	"sync/atomic"
    22  	"time"
    23  
    24  	"github.com/dustin/go-humanize"
    25  
    26  	"github.com/dgraph-io/dgraph/x"
    27  	"github.com/dgraph-io/ristretto/z"
    28  )
    29  
    30  type phase int32
    31  
    32  const (
    33  	nothing phase = iota
    34  	mapPhase
    35  	reducePhase
    36  )
    37  
    38  type progress struct {
    39  	nquadCount      int64
    40  	errCount        int64
    41  	mapEdgeCount    int64
    42  	reduceEdgeCount int64
    43  	reduceKeyCount  int64
    44  	numEncoding     int64
    45  
    46  	start       time.Time
    47  	startReduce time.Time
    48  
    49  	// shutdown is a bidirectional channel used to manage the stopping of the
    50  	// report goroutine. It handles both the request to stop the report
    51  	// goroutine, as well as the message back to say that the goroutine has
    52  	// stopped. The channel MUST be unbuffered for this to work.
    53  	shutdown chan struct{}
    54  
    55  	phase phase
    56  }
    57  
    58  func newProgress() *progress {
    59  	return &progress{
    60  		start:    time.Now(),
    61  		shutdown: make(chan struct{}),
    62  	}
    63  }
    64  
    65  func (p *progress) setPhase(ph phase) {
    66  	atomic.StoreInt32((*int32)(&p.phase), int32(ph))
    67  }
    68  
    69  func (p *progress) report() {
    70  	t := time.NewTicker(time.Second)
    71  	defer t.Stop()
    72  
    73  	z.StatsPrint() // Just print once.
    74  	for {
    75  		select {
    76  		case <-t.C:
    77  			p.reportOnce()
    78  		case <-p.shutdown:
    79  			p.shutdown <- struct{}{}
    80  			return
    81  		}
    82  	}
    83  }
    84  
    85  func (p *progress) reportOnce() {
    86  	mapEdgeCount := atomic.LoadInt64(&p.mapEdgeCount)
    87  	timestamp := time.Now().Format("15:04:05Z0700")
    88  
    89  	switch phase(atomic.LoadInt32((*int32)(&p.phase))) {
    90  	case nothing:
    91  	case mapPhase:
    92  		rdfCount := atomic.LoadInt64(&p.nquadCount)
    93  		errCount := atomic.LoadInt64(&p.errCount)
    94  		elapsed := time.Since(p.start)
    95  		fmt.Printf("[%s] MAP %s nquad_count:%s err_count:%s nquad_speed:%s/sec "+
    96  			"edge_count:%s edge_speed:%s/sec jemalloc: %s \n",
    97  			timestamp,
    98  			x.FixedDuration(elapsed),
    99  			niceFloat(float64(rdfCount)),
   100  			niceFloat(float64(errCount)),
   101  			niceFloat(float64(rdfCount)/elapsed.Seconds()),
   102  			niceFloat(float64(mapEdgeCount)),
   103  			niceFloat(float64(mapEdgeCount)/elapsed.Seconds()),
   104  			humanize.IBytes(uint64(z.NumAllocBytes())),
   105  		)
   106  	case reducePhase:
   107  		now := time.Now()
   108  		elapsed := time.Since(p.startReduce)
   109  		if p.startReduce.IsZero() {
   110  			p.startReduce = time.Now()
   111  			elapsed = time.Second
   112  		}
   113  		reduceKeyCount := atomic.LoadInt64(&p.reduceKeyCount)
   114  		reduceEdgeCount := atomic.LoadInt64(&p.reduceEdgeCount)
   115  		pct := ""
   116  		if mapEdgeCount != 0 {
   117  			pct = fmt.Sprintf("%.2f%% ", 100*float64(reduceEdgeCount)/float64(mapEdgeCount))
   118  		}
   119  		fmt.Printf("[%s] REDUCE %s %sedge_count:%s edge_speed:%s/sec "+
   120  			"plist_count:%s plist_speed:%s/sec. Num Encoding MBs: %d. jemalloc: %s \n",
   121  			timestamp,
   122  			x.FixedDuration(now.Sub(p.start)),
   123  			pct,
   124  			niceFloat(float64(reduceEdgeCount)),
   125  			niceFloat(float64(reduceEdgeCount)/elapsed.Seconds()),
   126  			niceFloat(float64(reduceKeyCount)),
   127  			niceFloat(float64(reduceKeyCount)/elapsed.Seconds()),
   128  			atomic.LoadInt64(&p.numEncoding)/(1<<20),
   129  			humanize.IBytes(uint64(z.NumAllocBytes())),
   130  		)
   131  	default:
   132  		x.AssertTruef(false, "invalid phase")
   133  	}
   134  }
   135  
   136  func (p *progress) endSummary() {
   137  	p.shutdown <- struct{}{}
   138  	<-p.shutdown
   139  
   140  	p.reportOnce()
   141  
   142  	total := x.FixedDuration(time.Since(p.start))
   143  	fmt.Printf("Total: %v\n", total)
   144  }
   145  
   146  var suffixes = [...]string{"", "k", "M", "G", "T"}
   147  
   148  func niceFloat(f float64) string {
   149  	idx := 0
   150  	for f >= 1000 {
   151  		f /= 1000
   152  		idx++
   153  	}
   154  	if idx >= len(suffixes) {
   155  		return fmt.Sprintf("%f", f)
   156  	}
   157  	suf := suffixes[idx]
   158  	switch {
   159  	case f >= 100:
   160  		return fmt.Sprintf("%.1f%s", f, suf)
   161  	case f >= 10:
   162  		return fmt.Sprintf("%.2f%s", f, suf)
   163  	default:
   164  		return fmt.Sprintf("%.3f%s", f, suf)
   165  	}
   166  }