(Blog) Goroutines demo

gowc.go

/*
  @file gowc.go

  Count the number of vowels present in a text file.
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
)

func main() {
	var (
		data       []byte
		num_cpus   int
		num_vowels uint64
		file_size  int64
		fileMap    map[int][]byte
	)

	/* Command line option. */
	var filename = flag.String("filename", "file.txt", "Name of the input file.")

	flag.Parse()

	file, err := os.Open(*filename)
	if err != nil {
		log.Fatal(err)
	}

	num_cpus = runtime.NumCPU()
	fileInfo, _ := file.Stat()
	file_size = fileInfo.Size()

	fmt.Println("Number of CPUs : ", num_cpus)

	fileMap = make(map[int][]byte)
	for i := 1; i <= num_cpus; i++ {
		/* FIXME : Exact partitioning logic omitted for simplicity. */
		data = make([]byte, file_size/int64(num_cpus))

		fragment_size, err := file.Read(data)
		fmt.Printf("#%d Fragment size : %d\n", i, fragment_size)

		if err != nil {
			log.Fatal(err)
		}

		fileMap[i] = append(fileMap[i], data...)
	}

	/* Create a buffered channel to get the result back from goroutines. */
	c := make(chan uint64, num_cpus)

	for i := 1; i <= num_cpus; i++ {
		go worker(fileMap[i], c)
	}

	/* Now, wait for all the goroutines to return. */
	for i := 1; i <= num_cpus; i++ {
		num_vowels += <-c // Wait for one goroutine to return.
	}

	fmt.Println("Total number of vowels : ", num_vowels)
}

func worker(data []byte, c chan uint64) {
	/* Write the result to the channel. */
	c <- countVowels(data)
}

func countVowels(data []byte) (count uint64) {
	for i := 0; i < len(data); i++ {
		switch data[i] {
		case 'a', 'A', 'e', 'E', 'i', 'I', 'o', 'O', 'u', 'U':
			count++
		}
	}
	return
}