go学习记录一:读取目录下的所有文件合并写入一个文件

77 阅读1分钟
package main

import (
   "bufio"
   "fmt"
   "io"
   "os"
   "strings"
   "sync/atomic"
)

var (
   fileChan = make(chan string)
   fc       = atomic.Int32{}
   exitChan = make(chan struct{})
)

func readFile(filename string) {
   file, err := os.Open(filename)
   if err != nil {
      fmt.Println(err)
   }
   defer file.Close()
   reader := bufio.NewReader(file)
   for {
      line, err := reader.ReadString('\n')
      line = strings.TrimSpace(line)
      if err == io.EOF {
         fmt.Println(fc)
         fc.Add(-1)
         break
      } else {
         fileChan <- line
      }
   }
}

func writeFile(filename string) {
   out, err := os.OpenFile(filename, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.ModePerm)
   if err != nil {
      fmt.Println(err.Error())
      return
   }
   defer out.Close()
   writer := bufio.NewWriter(out)
   for {
      select {
      case line := <-fileChan:
         writer.WriteString(line + "\n")
      default:
         break
      }
      if fc.Load() == 0 {
         break
      }
   }
   writer.Flush()
   exitChan <- struct{}{}
}

// g0
func main() {
   names, err := os.ReadDir("dir")
   if err != nil {
      fmt.Print(err)
   }
   paths := make([]string, 0)
   fc.Add(int32(len(names)))
   for _, name := range names {
      paths = append(paths, "dir/"+name.Name())
   }
   for _, name := range paths {
      go readFile(name)
   }
   go writeFile("merge")

   if _, ok := <-exitChan; ok {
      os.Exit(-1)
   }
}